Forums

from flask import Flask, request, jsonify, send_from_directory
import os
from torch import load, no_grad, device, max, topk  # Import only the required parts of torch
from torch.nn import Sequential, Linear, ReLU, Dropout  # Import specific layers
from torch.nn.functional import softmax  # Import softmax function
from torchvision.transforms import Compose, Resize, ToTensor, Normalize  # Import specific transforms
from torchvision.models import mobilenet_v2  # Import only the required model
from PIL import Image
import sqlite3

app = Flask(__name__)

# Directory configurations
UPLOAD_FOLDER = 'uploads'
MODEL_PATH = 'AppModel_with_classes_Epoch_150_Epoch_15.pt'
DB_PATH = 'missing_cats.db'
CAT_IMAGES_DIR = 'https://atelierhenger.nl/APP/cat_images/'
INDEX_HTML_PATH = '.'

# Ensure required directories exist
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)
if not os.path.exists(CAT_IMAGES_DIR):
    print(f"Warning: {CAT_IMAGES_DIR} directory does not exist. Please create it and add the images.")

# Load the model
def load_model():
    model = mobilenet_v2(weights=None)

    # Extract the number of classes from the checkpoint
    checkpoint = load(MODEL_PATH, map_location=device('cpu'))  # Use torch.load
    num_cat_classes = len(checkpoint['class_to_idx'])
    num_ras_classes = checkpoint.get('num_ras_classes', 0)  # Default to 0 if missing
    num_color_classes = checkpoint.get('num_color_classes', 0)  # Default to 0 if missing
    total_output_classes = num_cat_classes + num_ras_classes + num_color_classes

    # Adjust the classifier layer dynamically
    model.classifier = Sequential(
        Linear(model.last_channel, 512),
        ReLU(),
        Dropout(0.5),
        Linear(512, total_output_classes)  # Dynamically set the output size
    )

    # Load the state_dict
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()

    # Add readable idx_to_ras and idx_to_color mappings
    checkpoint['idx_to_ras'] = {
        0: 'Bengaal',
        1: 'Boskat X',
        2: 'Brits korthaar',
        3: 'Brits korthaar X',
        4: 'Europese korthaar of huiskat',
        5: 'Heilige Birmaan',
        6: 'Kruising',
        7: 'Langharig',
        8: 'Maine Coon',
        9: 'Maine Coon X',
        10: 'Overige rassen',
        11: 'Perzische',
        12: 'Perzische X',
        13: 'Ragdoll',
        14: 'Scottish Fold',
        15: 'Siamees',
        16: 'Somali',
        17: 'Turkse Angora X',
        18: 'Turkse Van'
    }

    checkpoint['idx_to_color'] = {
        0: 'Beige',
        1: 'Beige met wit',
        2: 'Bruin',
        3: 'Colourpoint blue point',
        4: 'Colourpoint lilac point',
        5: 'Colourpoint red point',
        6: 'Colourpoint seal point',
        7: 'Cypers bruin-zwart (gestreept)',
        8: 'Cypers bruin-zwart (gestreept) met wit',
        9: 'Cypers grijs (gestreept) met rood en wit',
        10: 'Cypers grijs-zwart (gestreept)',
        11: 'Cypers grijs-zwart (gestreept) met wit',
        12: 'Cypers rood (gestreept)',
        13: 'Cypers rood (gestreept) met wit',
        14: 'Grijs',
        15: 'Grijs met wit',
        16: 'Lapje',
        17: 'Rood',
        18: 'Rood met wit',
        19: 'Schildpad',
        20: 'Schildpad met wit',
        21: 'Sealpoint',
        22: 'Shaded silver',
        23: 'Wildkleur (met wit)',
        24: 'Wit',
        25: 'Wit met bruin',
        26: 'Wit met cypers bruin',
        27: 'Wit met cypers grijs',
        28: 'Wit met cypers rood',
        29: 'Wit met grijs',
        30: 'Wit met rood',
        31: 'Wit met zwart',
        32: 'Zwart',
        33: 'Zwart met wit'
    }

    return model, checkpoint

model, checkpoint = load_model()

# Define transformations
transform = Compose([
    Resize((224, 224)),
    ToTensor(),
    Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# Fetch metadata from the database
def get_cat_metadata(cat_number):
    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        query = "SELECT * FROM cats WHERE cat_number = ?"
        cursor.execute(query, (cat_number,))
        result = cursor.fetchone()
        column_names = [description[0] for description in cursor.description]
        conn.close()

        if result:
            metadata = dict(zip(column_names, result))
            # Check for images inside the folder for the cat_number
            supported_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']
            base_url = f"https://atelierhenger.nl/APP/cat_images/{cat_number}/"

            # Dynamically check for the first valid image in the folder
            for ext in supported_extensions:
                potential_url = f"{base_url}1{ext}"  # Always use the first image (e.g., 1.jpg, 1.png)
                metadata['image_url'] = potential_url
                break  # Use the first valid image found

            return metadata
        else:
            return None
    except sqlite3.Error as e:
        return {"error": str(e)}

@app.route('/')
def home():
    return send_from_directory('.', 'index.html')

@app.route('/upload', methods=['POST', 'GET'])
def upload_file():
    if request.method == 'GET':
        return '''
        <h1>Upload Endpoint</h1>
        <p>Send a POST request with an image file to this endpoint to check for matches.</p>
        '''
    try:
        if 'image' not in request.files:
            return jsonify({'error': 'No file uploaded'}), 400

        file = request.files['image']
        if file.filename == '':
            return jsonify({'error': 'No selected file'}), 400

        file_path = os.path.join(UPLOAD_FOLDER, file.filename)
        file.save(file_path)

        image = Image.open(file_path).convert('RGB')
        input_tensor = transform(image).unsqueeze(0)

        with no_grad():
            output = model(input_tensor)
            num_cat_classes = len(checkpoint['class_to_idx'])
            num_ras_classes = checkpoint['num_ras_classes']
            num_color_classes = checkpoint['num_color_classes']

            outputs_cat = output[:, :num_cat_classes]
            outputs_ras = output[:, num_cat_classes:num_cat_classes + num_ras_classes]
            outputs_color = output[:, num_cat_classes + num_ras_classes:]

            ras_confidence, ras_class = max(softmax(outputs_ras, dim=1), dim=1)
            color_confidence, color_class = max(softmax(outputs_color, dim=1), dim=1)

            predicted_ras = checkpoint['idx_to_ras'][ras_class.item()]
            predicted_color = checkpoint['idx_to_color'][color_class.item()]

            matches = []
            softmax_cat = softmax(outputs_cat, dim=1)
            topk_values, topk_indices = topk(softmax_cat, k=15, dim=1)
            for idx, confidence in zip(topk_indices[0].tolist(), topk_values[0].tolist()):
                cat_number = list(checkpoint['class_to_idx'].keys())[list(checkpoint['class_to_idx'].values()).index(idx)]
                metadata = get_cat_metadata(cat_number)
                if metadata:
                    metadata['confidence'] = f"{confidence * 100:.2f}"
                    matches.append(metadata)

            filtered_matches = [match for match in matches if match['ras'] == predicted_ras]
            sorted_matches = sorted(filtered_matches, key=lambda x: float(x['confidence']), reverse=True)

        return jsonify({
            'predicted_ras': predicted_ras,
            'predicted_ras_confidence': f"{ras_confidence.item() * 100:.2f}",
            'predicted_color': predicted_color,
            'predicted_color_confidence': f"{color_confidence.item() * 100:.2f}",
            'matches': sorted_matches
        })

    except Exception as e:
        print(f"Error processing the request: {e}")
        return jsonify({'error': 'An error occurred while processing the image.', 'details': str(e)}), 500