from flask import Flask, request, render_template, jsonify, send_from_directory import os import torch import numpy as np import cv2 from segment_anything import sam_model_registry, SamPredictor from werkzeug.utils import secure_filename import warnings import json # Initialisation de Flask app = Flask( __name__, template_folder='templates', static_folder='static' ) app.config['UPLOAD_FOLDER'] = os.path.join('static', 'uploads') os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True) # Charger le modèle SAM MODEL_TYPE = "vit_b" MODEL_PATH = os.path.join('models', 'sam_vit_b_01ec64.pth') device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print("Chargement du modèle SAM...") try: state_dict = torch.load(MODEL_PATH, map_location="cpu", weights_only=True) except TypeError: with warnings.catch_warnings(): warnings.simplefilter("ignore", category=UserWarning) state_dict = torch.load(MODEL_PATH, map_location="cpu") # Initialiser et charger le modèle sam = sam_model_registry[MODEL_TYPE]() sam.load_state_dict(state_dict, strict=False) sam.to(device=device) predictor = SamPredictor(sam) print("Modèle SAM chargé avec succès!") # Fonction pour générer une couleur unique pour chaque classe def get_color_for_class(class_name): np.random.seed(hash(class_name) % (2**32)) return tuple(np.random.randint(0, 256, size=3).tolist()) # Convertir un masque en bounding box au format YOLOv5 def mask_to_yolo_bbox(mask): y_indices, x_indices = np.where(mask > 0) if len(x_indices) == 0 or len(y_indices) == 0: return None x_min, x_max = x_indices.min(), x_indices.max() y_min, y_max = y_indices.min(), y_indices.max() # YOLOv5 format: x_center, y_center, width, height (normalized) x_center = (x_min + x_max) / 2 y_center = (y_min + y_max) / 2 width = x_max - x_min height = y_max - y_min return x_center, y_center, width, height @app.route('/', methods=['GET', 'POST']) def index(): if request.method == 'POST': files = request.files.getlist('images') # Get multiple files if not files: return "Aucun fichier sélectionné", 400 filenames = [] for file in files: filename = secure_filename(file.filename) filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename) file.save(filepath) filenames.append(filename) return render_template('index.html', uploaded_images=filenames, all_annotated=False) # Pour l'affichage des images déjà téléchargées uploaded_images = os.listdir(app.config['UPLOAD_FOLDER']) return render_template('index.html', uploaded_images=uploaded_images, all_annotated=False) @app.route('/uploads/') def uploaded_file(filename): return send_from_directory(app.config['UPLOAD_FOLDER'], filename) @app.route('/segment', methods=['POST']) def segment(): data = request.get_json() print("Données reçues :", data) # Log pour vérifier les données envoyées par le frontend image_names = data.get('image_names') points = data.get('points') if not image_names or not points: return jsonify({'success': False, 'error': 'Données manquantes'}), 400 output = [] for image_name in image_names: image_path = os.path.join(app.config['UPLOAD_FOLDER'], image_name) if not os.path.exists(image_path): return jsonify({'success': False, 'error': f'Image {image_name} non trouvée'}), 404 # Créer un dossier pour sauvegarder les résultats output_dir = os.path.join(app.config['UPLOAD_FOLDER'], os.path.splitext(image_name)[0]) os.makedirs(output_dir, exist_ok=True) # Charger l'image et effectuer la segmentation image = cv2.imread(image_path) image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) predictor.set_image(image_rgb) annotated_image = image.copy() # YOLOv5 annotation yolo_annotations = [] for point in points: x, y = point['x'], point['y'] class_name = point.get('class', 'Unknown') class_id = hash(class_name) % 1000 # Générer un ID unique basé sur le nom color = get_color_for_class(class_name) # Couleur unique pour chaque classe masks, _, _ = predictor.predict( point_coords=np.array([[x, y]]), point_labels=np.array([1]), multimask_output=False ) mask = masks[0] annotated_image[mask > 0] = color # Superposer le masque avec la couleur # Convertir le masque en bounding box YOLOv5 bbox = mask_to_yolo_bbox(mask) if bbox: x_center, y_center, width, height = bbox # Normaliser les valeurs x_center /= image.shape[1] y_center /= image.shape[0] width /= image.shape[1] height /= image.shape[0] yolo_annotations.append(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}") # Ajouter le texte de la classe cv2.putText(annotated_image, class_name, (int(x), int(y)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) # Texte blanc # Sauvegarder les résultats annotated_filename = f"annotated_{image_name}" annotated_path = os.path.join(output_dir, annotated_filename) cv2.imwrite(annotated_path, annotated_image) # Sauvegarder les annotations YOLOv5 yolo_path = os.path.join(output_dir, f"{os.path.splitext(image_name)[0]}.txt") with open(yolo_path, "w") as f: f.write("\n".join(yolo_annotations)) # Copier l'image originale dans le dossier original_copy_path = os.path.join(output_dir, image_name) if not os.path.exists(original_copy_path): os.rename(image_path, original_copy_path) # Renvoyer le chemin relatif pour affichage relative_output_dir = output_dir.replace("static/", "") output.append({ 'success': True, 'image': f"{relative_output_dir}/{annotated_filename}", 'yolo_annotations': f"{relative_output_dir}/{os.path.splitext(image_name)[0]}.txt" }) return jsonify(output) if __name__ == '__main__': app.run(debug=True, host='0.0.0.0', port=5000)