Upload app.py.bak

app.py.bak

@@ -0,0 +1,190 @@
+from typing import Tuple, Dict
+import gradio as gr
+import supervision as sv
+import numpy as np
+import cv2
+from huggingface_hub import hf_hub_download
+from ultralytics import YOLO
+
+# Define models
+MODEL_OPTIONS = {
+    "YOLOv11-Small": "medieval-yolo11s-seg.pt"
+}
+
+# Dictionary to store loaded models
+models: Dict[str, YOLO] = {}
+
+# Load all models
+for name, model_file in MODEL_OPTIONS.items():
+    try:
+        model_path = hf_hub_download(
+            repo_id="johnlockejrr/medieval-manuscript-yolov11-seg",
+            filename=model_file
+        )
+        models[name] = YOLO(model_path)
+    except Exception as e:
+        print(f"Error loading model {name}: {str(e)}")
+
+# Create annotators
+LABEL_ANNOTATOR = sv.LabelAnnotator(text_color=sv.Color.BLACK)
+MASK_ANNOTATOR = sv.MaskAnnotator()
+
+def process_masks(masks: np.ndarray, target_shape: Tuple[int, int]) -> np.ndarray:
+    """Process and resize masks to target shape"""
+    if masks is None:
+        return None
+        
+    processed_masks = []
+    h, w = target_shape
+    for mask in masks:
+        # Resize mask to target dimensions
+        resized_mask = cv2.resize(mask.astype(float), (w, h), interpolation=cv2.INTER_LINEAR)
+        # Threshold to create binary mask
+        processed_masks.append(resized_mask > 0.5)
+    
+    return np.array(processed_masks)
+
+def detect_and_annotate(
+    image: np.ndarray,
+    model_name: str,
+    conf_threshold: float,
+    iou_threshold: float
+) -> np.ndarray:
+    try:
+        if image is None:
+            return None
+            
+        model = models.get(model_name)
+        if model is None:
+            raise ValueError(f"Model {model_name} not loaded")
+        
+        # Perform inference
+        results = model.predict(
+            image,
+            conf=conf_threshold,
+            iou=iou_threshold
+        )[0]
+        
+        # Convert results to supervision Detections
+        boxes = results.boxes.xyxy.cpu().numpy()
+        confidence = results.boxes.conf.cpu().numpy()
+        class_ids = results.boxes.cls.cpu().numpy().astype(int)
+        
+        # Process masks
+        masks = None
+        if results.masks is not None:
+            masks = results.masks.data.cpu().numpy()
+            print(f"Original mask shape: {masks.shape}")  # Debug
+            
+            # Fix the shape mismatch - should be (num_masks, H, W)
+            if masks.shape[0] != len(boxes):
+                masks = np.transpose(masks, (2, 0, 1))  # Convert from (H,W,N) to (N,H,W)
+            
+            print(f"Processed mask shape: {masks.shape}")  # Debug
+            
+            # Resize masks to original image dimensions
+            h, w = image.shape[:2]
+            resized_masks = []
+            for mask in masks:
+                resized_mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_LINEAR)
+                resized_masks.append(resized_mask)
+            masks = np.array(resized_masks)
+            masks = masks > 0.5  # Convert to boolean
+            
+        # Create Detections object
+        detections = sv.Detections(
+            xyxy=boxes,
+            confidence=confidence,
+            class_id=class_ids,
+            mask=masks
+        )
+        
+        # Create labels with confidence scores
+        labels = [
+            f"{results.names[class_id]} ({conf:.2f})"
+            for class_id, conf in zip(class_ids, confidence)
+        ]
+
+        # Annotate image
+        annotated_image = image.copy()
+        if masks is not None:
+            annotated_image = MASK_ANNOTATOR.annotate(
+                scene=annotated_image, 
+                detections=detections
+            )
+        annotated_image = LABEL_ANNOTATOR.annotate(
+            scene=annotated_image,
+            detections=detections,
+            labels=labels
+        )
+        
+        return annotated_image
+        
+    except Exception as e:
+        print(f"Error during detection: {str(e)}")
+        return image
+
+# Create Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Medieval Manuscript Segmentation with YOLO")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Input Image", type='numpy')
+            with gr.Accordion("Detection Settings", open=True):
+                model_selector = gr.Dropdown(
+                    choices=list(MODEL_OPTIONS.keys()),
+                    value=list(MODEL_OPTIONS.keys())[0],
+                    label="Model"
+                )
+                conf_threshold = gr.Slider(
+                    label="Confidence Threshold",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.05,
+                    value=0.25
+                )
+                iou_threshold = gr.Slider(
+                    label="IoU Threshold",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.05,
+                    value=0.45
+                )
+            detect_btn = gr.Button("Detect", variant="primary")
+            clear_btn = gr.Button("Clear")
+                
+        with gr.Column():
+            output_image = gr.Image(label="Segmentation Result", type='numpy')
+
+    def process_image(image, model_name, conf_threshold, iou_threshold):
+        try:
+            if image is None:
+                return None, None
+            annotated_image = detect_and_annotate(image, model_name, conf_threshold, iou_threshold)
+            return image, annotated_image
+        except Exception as e:
+            print(f"Error in process_image: {str(e)}")
+            return image, image  # Fallback to original image
+
+    def clear():
+        return None, None
+
+    detect_btn.click(
+        process_image,
+        inputs=[input_image, model_selector, conf_threshold, iou_threshold],
+        outputs=[input_image, output_image]
+    )
+    clear_btn.click(
+        clear,
+        inputs=None,
+        outputs=[input_image, output_image]
+    )
+
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+        debug=True
+    )