added polygon support

comic_panel_extractor/annorator_server.py

@@ -8,6 +8,7 @@ import base64
 from io import BytesIO
 import shutil
 from .config import Config
+from typing import List, Optional, Union, Dict, Any
 
 app = APIRouter()
 
@@ -19,12 +20,21 @@ IMAGE_LABEL_ROOT = os.path.join(Config.current_path, "image_labels")
 CLASS_ID = 0
 
 # === Pydantic Models ===
+class Point(BaseModel):
+    x: float
+    y: float
+
 class Box(BaseModel):
-    left: int
-    top: int
-    width: int
-    height: int
-    type: str = "rect"
+    type: str = "bbox"  # "bbox" or "segmentation"
+    # For bbox
+    left: Optional[int] = None
+    top: Optional[int] = None
+    width: Optional[int] = None
+    height: Optional[int] = None
+    # For segmentation
+    points: Optional[List[Point]] = None
+    # Common fields
+    classId: int = CLASS_ID
     stroke: str = "#00ff00"
     strokeWidth: int = 3
     fill: str = "rgba(0, 255, 0, 0.2)"
@@ -32,11 +42,11 @@ class Box(BaseModel):
 
     @field_validator("left", "top", "width", "height", mode="before")
     def round_floats(cls, v):
-        return round(v)
+        return round(v) if v is not None else None
 
 class SaveAnnotationsRequest(BaseModel):
-    boxes: List[Box]
-    image_name: str  # Relative path like train/image1.jpg
+    annotations: List[Box]  # Changed from 'boxes' to 'annotations'
+    image_name: str
     original_width: int
     original_height: int
 
@@ -54,65 +64,162 @@ def get_label_path(image_name: str) -> str:
     return os.path.join(LABEL_ROOT, os.path.splitext(image_name)[0] + ".txt")
 
 # === Core Functions ===
-def load_yolo_boxes(image_path: str, label_path: str, detect: bool = False):
+def load_yolo_annotations(image_path: str, label_path: str, detect: bool = False):
+    """Load both bbox and segmentation annotations from YOLO format"""
     try:
         img = Image.open(image_path)
         w, h = img.size
-        boxes = []
+        annotations = []
+
+        # Auto-detect if needed
         if detect and not os.path.exists(label_path):
             from .yolo_manager import YOLOManager
             with YOLOManager() as yolo_manager:
                 weights_path = f'{Config.current_path}/{Config.YOLO_MODEL_NAME}.pt'
-
                 yolo_manager.load_model(weights_path)
-
-                # Run inference
-                _, label_path = yolo_manager.annotate_images(image_paths=[image_path], output_dir=IMAGE_LABEL_ROOT, save_image=False, label_path=label_path)
+                _, label_path = yolo_manager.annotate_images(
+                    image_paths=[image_path],
+                    output_dir=IMAGE_LABEL_ROOT,
+                    save_image=False,
+                    label_path=label_path
+                )
 
         if os.path.exists(label_path):
             with open(label_path, "r") as f:
                 for line in f:
                     parts = list(map(float, line.strip().split()))
-                    if len(parts) != 5:
+                    if len(parts) < 5:
                         continue
-                    _, xc, yc, bw, bh = parts
-                    left = int((xc - bw / 2) * w)
-                    top = int((yc - bh / 2) * h)
-                    width = int(bw * w)
-                    height = int(bh * h)
-                    boxes.append({
-                        "type": "rect",
-                        "left": left,
-                        "top": top,
-                        "width": width,
-                        "height": height,
-                        "stroke": "#00ff00",
-                        "strokeWidth": 3,
-                        "fill": "rgba(0, 255, 0, 0.2)",
-                        "saved": True
-                    })
-        return boxes, (w, h)
+
+                    class_id = int(parts[0])
+
+                    if len(parts) == 5:  # Bounding box format
+                        _, xc, yc, bw, bh = parts
+                        left = int((xc - bw / 2) * w)
+                        top = int((yc - bh / 2) * h)
+                        width = int(bw * w)
+                        height = int(bh * h)
+
+                        annotations.append({
+                            "type": "bbox",
+                            "left": left,
+                            "top": top,
+                            "width": width,
+                            "height": height,
+                            "classId": class_id,
+                            "stroke": "#00ff00",
+                            "strokeWidth": 3,
+                            "fill": "rgba(0, 255, 0, 0.2)",
+                            "saved": True
+                        })
+
+                    elif len(parts) > 5 and len(parts) % 2 == 1:  # Segmentation format
+                        # Skip class_id, then pairs of x,y coordinates
+                        coords = parts[1:]
+                        if len(coords) >= 6:  # At least 3 points
+                            points = []
+                            for i in range(0, len(coords), 2):
+                                if i + 1 < len(coords):
+                                    x = coords[i] * w
+                                    y = coords[i + 1] * h
+                                    points.append({"x": x, "y": y})
+
+                            annotations.append({
+                                "type": "segmentation",
+                                "points": points,
+                                "classId": class_id,
+                                "stroke": "#00ff00",
+                                "strokeWidth": 3,
+                                "fill": "rgba(0, 255, 0, 0.2)",
+                                "saved": True
+                            })
+
+        return annotations, (w, h)
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error loading data: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error loading annotations: {str(e)}")
 
-def save_yolo_annotations(boxes: List[Box], original_size: tuple, label_path: str):
+def save_yolo_annotations(annotations: List[Box], original_size: tuple, label_path: str):
+    """Save annotations in YOLO format (both bbox and segmentation)"""
     os.makedirs(os.path.dirname(label_path), exist_ok=True)
     w, h = original_size
+
     try:
         with open(label_path, "w") as f:
-            for box in boxes:
-                left, top, width, height = box.left, box.top, box.width, box.height
-                xc = (left + width / 2) / w
-                yc = (top + height / 2) / h
-                bw = width / w
-                bh = height / h
-                f.write(f"{CLASS_ID} {xc:.6f} {yc:.6f} {bw:.6f} {bh:.6f}\n")
-
+            # Generate YOLO format from annotations
+            for annotation in annotations:
+                if annotation.type == "bbox":
+                    left, top, width, height = annotation.left, annotation.top, annotation.width, annotation.height
+                    xc = (left + width / 2) / w
+                    yc = (top + height / 2) / h
+                    bw = width / w
+                    bh = height / h
+                    f.write(f"{annotation.classId} {xc:.6f} {yc:.6f} {bw:.6f} {bh:.6f}\n")
+
+                elif annotation.type == "segmentation" and annotation.points:
+                    # Convert points to normalized coordinates
+                    normalized_points = []
+                    for point in annotation.points:
+                        normalized_points.extend([point.x / w, point.y / h])
+
+                    coords_str = " ".join(f"{coord:.6f}" for coord in normalized_points)
+                    f.write(f"{annotation.classId} {coords_str}\n")
+
+        # Copy to image_labels directory
         shutil.copy2(label_path, f"{IMAGE_LABEL_ROOT}/{os.path.basename(label_path)}")
         return True
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error saving annotations: {str(e)}")
 
+def parse_yolo_line(line: str, image_width: int, image_height: int) -> Dict[str, Any]:
+    """Parse a single YOLO format line and return annotation dict"""
+    parts = list(map(float, line.strip().split()))
+    if len(parts) < 5:
+        return None
+
+    class_id = int(parts[0])
+
+    if len(parts) == 5:  # Bounding box
+        _, xc, yc, bw, bh = parts
+        left = int((xc - bw / 2) * image_width)
+        top = int((yc - bh / 2) * image_height)
+        width = int(bw * image_width)
+        height = int(bh * image_height)
+
+        return {
+            "type": "bbox",
+            "left": left,
+            "top": top,
+            "width": width,
+            "height": height,
+            "classId": class_id,
+            "stroke": "#00ff00",
+            "strokeWidth": 3,
+            "fill": "rgba(0, 255, 0, 0.2)",
+            "saved": True
+        }
+
+    elif len(parts) > 5 and len(parts) % 2 == 1:  # Segmentation
+        coords = parts[1:]
+        if len(coords) >= 6:  # At least 3 points
+            points = []
+            for i in range(0, len(coords), 2):
+                if i + 1 < len(coords):
+                    x = coords[i] * image_width
+                    y = coords[i + 1] * image_height
+                    points.append({"x": x, "y": y})
+
+            return {
+                "type": "segmentation",
+                "points": points,
+                "classId": class_id,
+                "stroke": "#00ff00",
+                "strokeWidth": 3,
+                "fill": "rgba(0, 255, 0, 0.2)",
+                "saved": True
+            }
+
+    return None
+
 # === API Routes ===
 
 @app.get("/api/annotate/images", response_model=List[ImageInfo])
@@ -162,24 +269,24 @@ async def get_annotations(image_name: str):
     if not os.path.exists(image_path):
         raise HTTPException(status_code=404, detail="Image not found")
 
-    boxes, (width, height) = load_yolo_boxes(image_path, label_path)
+    annotations, (width, height) = load_yolo_annotations(image_path, label_path)
     return {
-        "boxes": boxes,
+        "annotations": annotations,  # Changed from "boxes"
         "original_width": width,
         "original_height": height
     }
 
 @app.get("/api/annotate/detect_annotations/{image_name:path}")
-async def get_annotations(image_name: str):
+async def get_detected_annotations(image_name: str):
     image_path = get_image_path(image_name)
     label_path = get_label_path(image_name)
 
     if not os.path.exists(image_path):
         raise HTTPException(status_code=404, detail="Image not found")
 
-    boxes, (width, height) = load_yolo_boxes(image_path, label_path, True)
+    annotations, (width, height) = load_yolo_annotations(image_path, label_path, True)
     return {
-        "boxes": boxes,
+        "annotations": annotations,
         "original_width": width,
         "original_height": height
     }
@@ -188,11 +295,11 @@ async def get_annotations(image_name: str):
 async def save_annotations(request: SaveAnnotationsRequest):
     label_path = get_label_path(request.image_name)
     success = save_yolo_annotations(
-        request.boxes,
+        request.annotations,
         (request.original_width, request.original_height),
         label_path
     )
-    return {"message": f"Saved {len(request.boxes)} annotations successfully"}
+    return {"message": f"Saved {len(request.annotations)} annotations successfully"}
 
 @app.delete("/api/annotate/annotations/{image_name:path}")
 async def delete_annotations(image_name: str):

comic_panel_extractor/config.py

@@ -1,6 +1,9 @@
 from dataclasses import dataclass
 import os
 
+from dotenv import load_dotenv
+load_dotenv()
+
 @dataclass
 class Config:
 	"""Configuration settings for the comic-to-video pipeline."""

comic_panel_extractor/static/annotator.html

@@ -499,6 +499,21 @@
                         Next →
                     </button>
                 </div>
+                <!-- Annotation Mode -->
+                <div class="sidebar-section">
+                    <div class="section-title">Annotation Mode</div>
+                    <div class="form-field">
+                        <select class="form-select" id="annotationMode">
+                            <option value="bbox">Bounding Box (YOLO)</option>
+                            <option value="segmentation">Segmentation (YOLO-seg)</option>
+                        </select>
+                    </div>
+                    <div class="form-field">
+                        <label class="form-label">Class ID</label>
+                        <input type="number" class="form-input" id="classId" value="0" min="0" max="100">
+                    </div>
+                </div>
+
 
                 <div class="file-upload">
                     <input type="file" id="uploadFile" accept="image/*">
@@ -527,7 +542,7 @@
                 <div class="section-title">Current Image</div>
                 <div class="info-card">
                     <div class="info-row">
-                        <span class="info-label">Boxes</span>
+                        <span class="info-label">Annotations</span>
                         <span class="info-value" id="boxCount">0</span>
                     </div>
                     <div class="info-row">
@@ -582,7 +597,7 @@
         <div class="canvas-area">
             <div class="canvas-toolbar">
                 <span id="file_name" style="font-size: 13px; color: #4a5568;">
-                    Click and drag to create annotation boxes • Select boxes to move or resize
+                    Click and drag to create annotation annotations • Select annotations to move or resize
                 </span>
             </div>
 
@@ -604,7 +619,7 @@
             constructor() {
                 this.canvas = document.getElementById('annotationCanvas');
                 this.ctx = this.canvas.getContext('2d');
-                this.boxes = [];
+                this.annotations = [];
                 this.images = [];
                 this.currentImageIndex = -1;
                 this.currentImage = null;
@@ -618,6 +633,12 @@
                 this.startY = 0;
                 this.currentBox = null;
 
+                // Segmentation state
+                this.annotationMode = 'bbox'; // 'bbox' or 'segmentation'
+                this.currentPolygon = null;
+                this.isDrawingPolygon = false;
+                this.polygonPoints = [];
+
                 // Box editing state
                 this.selectedBoxIndex = -1;
                 this.isDragging = false;
@@ -637,6 +658,16 @@
             }
 
             setupEventListeners() {
+                // Mode selection
+                document.getElementById('annotationMode').addEventListener('change', (e) => {
+                    this.annotationMode = e.target.value;
+                    this.selectedBoxIndex = -1;
+                    this.polygonPoints = [];
+                    this.isDrawingPolygon = false;
+                    this.updateCanvasCursor();
+                    this.drawCanvas();
+                });
+
                 // Image selection
                 document.getElementById('imageSelect').addEventListener('change', (e) => {
                     if (e.target.value) {
@@ -662,7 +693,7 @@
                 // Action buttons
                 document.getElementById('saveBtn').addEventListener('click', () => this.saveAnnotations());
                 document.getElementById('undoBtn').addEventListener('click', () => this.undoLastBox());
-                document.getElementById('clearBtn').addEventListener('click', () => this.clearAllBoxes());
+                document.getElementById('clearBtn').addEventListener('click', () => this.clearAllAnnotations());
                 document.getElementById('reloadBtn').addEventListener('click', () => this.reloadAnnotations());
                 document.getElementById('detectBtn').addEventListener('click', () => this.detectAnnotations());
                 document.getElementById('downloadBtn').addEventListener('click', () => this.downloadAnnotations());
@@ -672,6 +703,7 @@
                 this.canvas.addEventListener('mousemove', (e) => this.onMouseMove(e));
                 this.canvas.addEventListener('mouseup', () => this.onMouseUp());
                 this.canvas.addEventListener('mouseleave', () => this.onMouseUp());
+                this.canvas.addEventListener('dblclick', (e) => this.onDoubleClick(e));
 
                 // Keyboard events
                 document.addEventListener('keydown', (e) => this.onKeyDown(e));
@@ -680,6 +712,14 @@
                 this.canvas.tabIndex = 0;
             }
 
+            updateCanvasCursor() {
+                if (this.annotationMode === 'segmentation') {
+                    this.canvas.style.cursor = 'crosshair';
+                } else {
+                    this.canvas.style.cursor = 'crosshair';
+                }
+            }
+
             async loadImages() {
                 try {
                     const response = await fetch('/api/annotate/images');
@@ -757,7 +797,7 @@
                     this.currentImage = imageName;
                     this.originalWidth = imageData.width;
                     this.originalHeight = imageData.height;
-                    this.boxes = annotationsData.boxes || [];
+                    this.annotations = annotationsData.annotations || [];
                     this.selectedBoxIndex = -1;
 
                     // Load and draw image
@@ -779,7 +819,7 @@
 
                     // Update info panel
                     document.getElementById('currentImageInfo').style.display = 'block';
-                    document.getElementById('boxCount').textContent = this.boxes.length;
+                    document.getElementById('boxCount').textContent = this.annotations.length;
                     document.getElementById('imageSize').textContent = `${imageData.width}×${imageData.height}`;
                     document.getElementById('selectedBoxInfo').textContent = 'None';
 
@@ -793,29 +833,32 @@
             drawCanvas() {
                 this.ctx.clearRect(0, 0, this.canvas.width, this.canvas.height);
 
-                // Draw background image
                 if (this.backgroundImage) {
                     this.ctx.drawImage(this.backgroundImage, 0, 0);
                 }
 
-                // Draw existing boxes
-                this.boxes.forEach((box, index) => {
+                // Draw existing annotations
+                this.annotations.forEach((annotation, index) => {
                     let strokeColor = '#00ff00';
                     let fillColor = 'rgba(0, 255, 0, 0.2)';
 
                     if (index === this.selectedBoxIndex) {
                         strokeColor = '#0066ff';
                         fillColor = 'rgba(0, 102, 255, 0.3)';
-                    } else if (!box.saved) {
+                    } else if (!annotation.saved) {
                         strokeColor = '#ff0000';
                         fillColor = 'rgba(255, 0, 0, 0.2)';
                     }
 
-                    this.drawBox(box.left, box.top, box.width, box.height, strokeColor, fillColor);
+                    if (annotation.type === 'segmentation') {
+                        this.drawPolygon(annotation.points, strokeColor, fillColor);
+                    } else {
+                        this.drawBox(annotation.left, annotation.top, annotation.width, annotation.height, strokeColor, fillColor);
+                    }
 
-                    // Draw resize handles for selected box
-                    if (index === this.selectedBoxIndex) {
-                        this.drawResizeHandles(box);
+                    // Draw resize handles for selected bbox
+                    if (index === this.selectedBoxIndex && annotation.type === 'bbox') {
+                        this.drawResizeHandles(annotation);
                     }
                 });
 
@@ -824,8 +867,173 @@
                     this.drawBox(this.currentBox.left, this.currentBox.top,
                         this.currentBox.width, this.currentBox.height, '#ff0000', 'rgba(255, 0, 0, 0.3)');
                 }
+
+                // Draw current polygon being drawn
+                if (this.currentPolygon && this.currentPolygon.points.length > 0) {
+                    this.drawPolygon(this.currentPolygon.points, '#ff0000', 'rgba(255, 0, 0, 0.3)');
+                }
+            }
+            drawPolygon(points, strokeColor, fillColor, isSelected = false) {
+                if (points.length < 2) return;
+
+                this.ctx.strokeStyle = strokeColor;
+                this.ctx.fillStyle = fillColor;
+                this.ctx.lineWidth = 3;
+
+                this.ctx.beginPath();
+                this.ctx.moveTo(points[0].x, points[0].y);
+
+                for (let i = 1; i < points.length; i++) {
+                    this.ctx.lineTo(points[i].x, points[i].y);
+                }
+
+                if (points.length > 2) {
+                    this.ctx.closePath();
+                    this.ctx.fill();
+                }
+                this.ctx.stroke();
+
+                // Draw points (larger for selected polygon)
+                const pointRadius = isSelected ? 6 : 4;
+                points.forEach((point, index) => {
+                    this.ctx.fillStyle = isSelected ? '#0066ff' : strokeColor;
+                    this.ctx.beginPath();
+                    this.ctx.arc(point.x, point.y, pointRadius, 0, 2 * Math.PI);
+                    this.ctx.fill();
+
+                    // Draw point outline for selected polygon
+                    if (isSelected) {
+                        this.ctx.strokeStyle = '#ffffff';
+                        this.ctx.lineWidth = 2;
+                        this.ctx.stroke();
+                    }
+                });
+
+                // Draw edge midpoints for selected polygon (for adding points)
+                if (isSelected && points.length > 2) {
+                    this.ctx.fillStyle = 'rgba(0, 102, 255, 0.7)';
+                    for (let i = 0; i < points.length; i++) {
+                        const nextIndex = (i + 1) % points.length;
+                        const midX = (points[i].x + points[nextIndex].x) / 2;
+                        const midY = (points[i].y + points[nextIndex].y) / 2;
+
+                        this.ctx.beginPath();
+                        this.ctx.arc(midX, midY, 3, 0, 2 * Math.PI);
+                        this.ctx.fill();
+                    }
+                }
+            }
+
+            // Polygon-specific editing methods
+            getPolygonAtPosition(x, y) {
+                for (let i = this.annotations.length - 1; i >= 0; i--) {
+                    const annotation = this.annotations[i];
+                    if (annotation.type === 'segmentation' && this.isPointInPolygon(x, y, annotation.points)) {
+                        return i;
+                    }
+                }
+                return -1;
             }
 
+            isPointInPolygon(x, y, points) {
+                let inside = false;
+                for (let i = 0, j = points.length - 1; i < points.length; j = i++) {
+                    if (((points[i].y > y) !== (points[j].y > y)) &&
+                        (x < (points[j].x - points[i].x) * (y - points[i].y) / (points[j].y - points[i].y) + points[i].x)) {
+                        inside = !inside;
+                    }
+                }
+                return inside;
+            }
+
+            getPolygonPointAtPosition(x, y, tolerance = 8) {
+                if (this.selectedBoxIndex >= 0 && this.annotations[this.selectedBoxIndex].type === 'segmentation') {
+                    const points = this.annotations[this.selectedBoxIndex].points;
+                    for (let i = 0; i < points.length; i++) {
+                        const distance = Math.sqrt(Math.pow(x - points[i].x, 2) + Math.pow(y - points[i].y, 2));
+                        if (distance <= tolerance) {
+                            return i;
+                        }
+                    }
+                }
+                return -1;
+            }
+
+            getPolygonEdgeAtPosition(x, y, tolerance = 5) {
+                if (this.selectedBoxIndex >= 0 && this.annotations[this.selectedBoxIndex].type === 'segmentation') {
+                    const points = this.annotations[this.selectedBoxIndex].points;
+                    for (let i = 0; i < points.length; i++) {
+                        const nextIndex = (i + 1) % points.length;
+                        const p1 = points[i];
+                        const p2 = points[nextIndex];
+
+                        const distance = this.distanceToLineSegment(x, y, p1.x, p1.y, p2.x, p2.y);
+                        if (distance <= tolerance) {
+                            return i; // Return the index of the first point of the edge
+                        }
+                    }
+                }
+                return -1;
+            }
+
+            distanceToLineSegment(px, py, x1, y1, x2, y2) {
+                const dx = x2 - x1;
+                const dy = y2 - y1;
+                const length = Math.sqrt(dx * dx + dy * dy);
+
+                if (length === 0) return Math.sqrt((px - x1) ** 2 + (py - y1) ** 2);
+
+                const t = Math.max(0, Math.min(1, ((px - x1) * dx + (py - y1) * dy) / (length * length)));
+                const projectionX = x1 + t * dx;
+                const projectionY = y1 + t * dy;
+
+                return Math.sqrt((px - projectionX) ** 2 + (py - projectionY) ** 2);
+            }
+
+            movePolygon(index, deltaX, deltaY) {
+                if (index >= 0 && index < this.annotations.length && this.annotations[index].type === 'segmentation') {
+                    const annotation = this.annotations[index];
+                    annotation.points = annotation.points.map(point => ({
+                        x: Math.max(0, Math.min(this.originalWidth, point.x + deltaX)),
+                        y: Math.max(0, Math.min(this.originalHeight, point.y + deltaY))
+                    }));
+                    annotation.saved = false;
+                }
+            }
+
+            movePolygonPoint(annotationIndex, pointIndex, newX, newY) {
+                if (annotationIndex >= 0 && annotationIndex < this.annotations.length) {
+                    const annotation = this.annotations[annotationIndex];
+                    if (annotation.type === 'segmentation' && pointIndex >= 0 && pointIndex < annotation.points.length) {
+                        annotation.points[pointIndex].x = Math.max(0, Math.min(this.originalWidth, newX));
+                        annotation.points[pointIndex].y = Math.max(0, Math.min(this.originalHeight, newY));
+                        annotation.saved = false;
+                    }
+                }
+            }
+
+            addPolygonPoint(annotationIndex, edgeIndex, x, y) {
+                if (annotationIndex >= 0 && annotationIndex < this.annotations.length) {
+                    const annotation = this.annotations[annotationIndex];
+                    if (annotation.type === 'segmentation') {
+                        const insertIndex = edgeIndex + 1;
+                        annotation.points.splice(insertIndex, 0, { x, y });
+                        annotation.saved = false;
+                    }
+                }
+            }
+
+            removePolygonPoint(annotationIndex, pointIndex) {
+                if (annotationIndex >= 0 && annotationIndex < this.annotations.length) {
+                    const annotation = this.annotations[annotationIndex];
+                    if (annotation.type === 'segmentation' && annotation.points.length > 3) {
+                        annotation.points.splice(pointIndex, 1);
+                        annotation.saved = false;
+                    }
+                }
+            }
+
+
             drawBox(x, y, width, height, strokeColor, fillColor) {
                 this.ctx.strokeStyle = strokeColor;
                 this.ctx.fillStyle = fillColor;
@@ -864,7 +1072,59 @@
                 this.lastMouseX = pos.x;
                 this.lastMouseY = pos.y;
 
-                // Check if clicking on a resize handle
+                if (this.annotationMode === 'segmentation') {
+                    // Check if clicking on a point of selected polygon
+                    if (this.selectedBoxIndex >= 0) {
+                        const pointIndex = this.getPolygonPointAtPosition(pos.x, pos.y);
+                        if (pointIndex >= 0) {
+                            this.isDraggingPoint = true;
+                            this.draggingPointIndex = pointIndex;
+                            this.canvas.style.cursor = 'move';
+                            return;
+                        }
+
+                        // Check if clicking on an edge to add a point
+                        const edgeIndex = this.getPolygonEdgeAtPosition(pos.x, pos.y);
+                        if (edgeIndex >= 0 && e.ctrlKey) { // Ctrl+click to add point
+                            this.addPolygonPoint(this.selectedBoxIndex, edgeIndex, pos.x, pos.y);
+                            this.drawCanvas();
+                            return;
+                        }
+                    }
+
+                    // Check if clicking on a polygon
+                    const clickedPolygonIndex = this.getPolygonAtPosition(pos.x, pos.y);
+                    if (clickedPolygonIndex >= 0) {
+                        this.selectedBoxIndex = clickedPolygonIndex;
+                        this.isDragging = true;
+                        this.dragStartX = pos.x;
+                        this.dragStartY = pos.y;
+                        this.canvas.style.cursor = 'move';
+                        this.updateSelectedBoxInfo();
+                        this.drawCanvas();
+                        return;
+                    }
+
+                    // Start new polygon if not drawing one
+                    if (!this.isDrawingPolygon) {
+                        this.selectedBoxIndex = -1;
+                        this.isDrawingPolygon = true;
+                        this.polygonPoints = [pos];
+                        this.currentPolygon = {
+                            points: [...this.polygonPoints],
+                            classId: parseInt(document.getElementById('classId').value) || 0,
+                            saved: false
+                        };
+                    } else {
+                        // Add point to current polygon
+                        this.polygonPoints.push(pos);
+                        this.currentPolygon.points = [...this.polygonPoints];
+                    }
+                    this.drawCanvas();
+                    return;
+                }
+
+                // Original bbox logic remains the same...
                 if (this.selectedBoxIndex >= 0) {
                     const handle = this.getResizeHandle(pos.x, pos.y);
                     if (handle) {
@@ -875,7 +1135,6 @@
                     }
                 }
 
-                // Check if clicking on an existing box
                 const clickedBoxIndex = this.getBoxAtPosition(pos.x, pos.y);
                 if (clickedBoxIndex >= 0) {
                     this.selectedBoxIndex = clickedBoxIndex;
@@ -888,7 +1147,6 @@
                     return;
                 }
 
-                // Start drawing new box
                 this.selectedBoxIndex = -1;
                 this.startX = pos.x;
                 this.startY = pos.y;
@@ -897,11 +1155,73 @@
                 this.updateSelectedBoxInfo();
             }
 
+            onDoubleClick(e) {
+                if (this.annotationMode === 'segmentation' && this.isDrawingPolygon && this.polygonPoints.length >= 3) {
+                    // Finish polygon
+                    this.finishPolygon();
+                }
+            }
+
+            finishPolygon() {
+                if (this.polygonPoints.length >= 3) {
+                    const annotation = {
+                        type: 'segmentation',
+                        points: [...this.polygonPoints],
+                        classId: parseInt(document.getElementById('classId').value) || 0,
+                        saved: false
+                    };
+
+                    this.annotations.push(annotation);
+                    this.selectedBoxIndex = this.annotations.length - 1;
+                    document.getElementById('boxCount').textContent = this.annotations.length;
+                    this.updateSelectedBoxInfo();
+                }
+
+                this.isDrawingPolygon = false;
+                this.polygonPoints = [];
+                this.currentPolygon = null;
+                this.drawCanvas();
+            }
+
             onMouseMove(e) {
                 if (!this.currentImage) return;
 
                 const pos = this.getMousePos(e);
 
+                if (this.annotationMode === 'segmentation') {
+                    if (this.isDraggingPoint && this.selectedBoxIndex >= 0) {
+                        // Move individual point
+                        this.movePolygonPoint(this.selectedBoxIndex, this.draggingPointIndex, pos.x, pos.y);
+                        this.drawCanvas();
+                        return;
+                    }
+
+                    if (this.isDragging && this.selectedBoxIndex >= 0) {
+                        // Move entire polygon
+                        const deltaX = pos.x - this.lastMouseX;
+                        const deltaY = pos.y - this.lastMouseY;
+                        this.movePolygon(this.selectedBoxIndex, deltaX, deltaY);
+                        this.lastMouseX = pos.x;
+                        this.lastMouseY = pos.y;
+                        this.drawCanvas();
+                        return;
+                    }
+
+                    if (this.isDrawingPolygon) {
+                        // Update current polygon preview
+                        if (this.currentPolygon) {
+                            this.currentPolygon.points = [...this.polygonPoints, pos];
+                            this.drawCanvas();
+                        }
+                        return;
+                    }
+
+                    // Update cursor based on what's under mouse
+                    this.updatePolygonCursor(pos.x, pos.y);
+                    return;
+                }
+
+                // Original bbox mouse move logic...
                 if (this.isResizing && this.selectedBoxIndex >= 0) {
                     this.resizeBox(pos.x, pos.y);
                     this.drawCanvas();
@@ -917,43 +1237,46 @@
                     this.currentBox.height = pos.y - this.startY;
                     this.drawCanvas();
                 } else {
-                    // Update cursor based on what's under mouse
                     this.updateCursor(pos.x, pos.y);
                 }
             }
 
             onMouseUp() {
+                if (this.annotationMode === 'segmentation') {
+                    this.isDraggingPoint = false;
+                    this.draggingPointIndex = -1;
+                    this.isDragging = false;
+                    this.canvas.style.cursor = 'crosshair';
+                    return;
+                }
+
+                // Original bbox mouse up logic...
                 if (this.isDrawing && this.currentBox) {
-                    // Only add box if it has meaningful size
                     if (Math.abs(this.currentBox.width) > 10 && Math.abs(this.currentBox.height) > 10) {
                         let left = Math.min(this.startX, this.startX + this.currentBox.width);
                         let top = Math.min(this.startY, this.startY + this.currentBox.height);
                         let width = Math.abs(this.currentBox.width);
                         let height = Math.abs(this.currentBox.height);
 
-                        // Apply boundary conditions
                         left = Math.max(0, left);
                         top = Math.max(0, top);
                         width = Math.min(width, this.originalWidth - left);
                         height = Math.min(height, this.originalHeight - top);
 
-                        // Only create box if it still has valid dimensions after boundary check
                         if (width > 10 && height > 10) {
                             const box = {
+                                type: 'bbox',
                                 left: left,
                                 top: top,
                                 width: width,
                                 height: height,
-                                type: 'rect',
-                                stroke: '#ff0000',
-                                strokeWidth: 3,
-                                fill: 'rgba(255, 0, 0, 0.3)',
+                                classId: parseInt(document.getElementById('classId').value) || 0,
                                 saved: false
                             };
 
-                            this.boxes.push(box);
-                            this.selectedBoxIndex = this.boxes.length - 1;
-                            document.getElementById('boxCount').textContent = this.boxes.length;
+                            this.annotations.push(box);
+                            this.selectedBoxIndex = this.annotations.length - 1;
+                            document.getElementById('boxCount').textContent = this.annotations.length;
                             this.updateSelectedBoxInfo();
                         }
                     }
@@ -967,7 +1290,33 @@
                 this.canvas.style.cursor = 'crosshair';
                 this.drawCanvas();
             }
+            updatePolygonCursor(x, y) {
+                if (this.selectedBoxIndex >= 0 && this.annotations[this.selectedBoxIndex].type === 'segmentation') {
+                    const pointIndex = this.getPolygonPointAtPosition(x, y);
+                    if (pointIndex >= 0) {
+                        this.canvas.style.cursor = 'move';
+                        return;
+                    }
+
+                    const edgeIndex = this.getPolygonEdgeAtPosition(x, y);
+                    if (edgeIndex >= 0) {
+                        this.canvas.style.cursor = 'pointer';
+                        return;
+                    }
 
+                    if (this.isPointInPolygon(x, y, this.annotations[this.selectedBoxIndex].points)) {
+                        this.canvas.style.cursor = 'move';
+                        return;
+                    }
+                }
+
+                const polygonIndex = this.getPolygonAtPosition(x, y);
+                if (polygonIndex >= 0) {
+                    this.canvas.style.cursor = 'pointer';
+                } else {
+                    this.canvas.style.cursor = 'crosshair';
+                }
+            }
 
             updateCursor(x, y) {
                 if (this.selectedBoxIndex >= 0) {
@@ -989,7 +1338,7 @@
             getResizeHandle(x, y) {
                 if (this.selectedBoxIndex < 0) return null;
 
-                const box = this.boxes[this.selectedBoxIndex];
+                const box = this.annotations[this.selectedBoxIndex];
                 const handleSize = 8;
                 const tolerance = 5;
 
@@ -1013,8 +1362,8 @@
             }
 
             getBoxAtPosition(x, y) {
-                for (let i = this.boxes.length - 1; i >= 0; i--) {
-                    const box = this.boxes[i];
+                for (let i = this.annotations.length - 1; i >= 0; i--) {
+                    const box = this.annotations[i];
                     if (x >= box.left && x <= box.left + box.width &&
                         y >= box.top && y <= box.top + box.height) {
                         return i;
@@ -1026,7 +1375,7 @@
             resizeBox(x, y) {
                 if (this.selectedBoxIndex < 0 || !this.resizeHandle) return;
 
-                const box = this.boxes[this.selectedBoxIndex];
+                const box = this.annotations[this.selectedBoxIndex];
                 const handle = this.resizeHandle;
 
                 switch (handle.type) {
@@ -1119,9 +1468,9 @@
 
 
             moveBox(boxIndex, deltaX, deltaY) {
-                if (boxIndex < 0 || boxIndex >= this.boxes.length) return;
+                if (boxIndex < 0 || boxIndex >= this.annotations.length) return;
 
-                const box = this.boxes[boxIndex];
+                const box = this.annotations[boxIndex];
 
                 // Calculate new position
                 let newLeft = box.left + deltaX;
@@ -1212,7 +1561,7 @@
             resizeSelectedBox(deltaWidth, deltaHeight) {
                 if (this.selectedBoxIndex < 0) return;
 
-                const box = this.boxes[this.selectedBoxIndex];
+                const box = this.annotations[this.selectedBoxIndex];
 
                 // Calculate new dimensions
                 let newWidth = box.width + deltaWidth;
@@ -1250,9 +1599,9 @@
 
             deleteSelectedBox() {
                 if (this.selectedBoxIndex >= 0) {
-                    this.boxes.splice(this.selectedBoxIndex, 1);
+                    this.annotations.splice(this.selectedBoxIndex, 1);
                     this.selectedBoxIndex = -1;
-                    document.getElementById('boxCount').textContent = this.boxes.length;
+                    document.getElementById('boxCount').textContent = this.annotations.length;
                     this.updateSelectedBoxInfo();
                     this.drawCanvas();
                     // this.showAlert('Box deleted', 'info');
@@ -1262,7 +1611,7 @@
             updateSelectedBoxInfo() {
                 const selectedBoxInfo = document.getElementById('selectedBoxInfo');
                 if (this.selectedBoxIndex >= 0) {
-                    const box = this.boxes[this.selectedBoxIndex];
+                    const box = this.annotations[this.selectedBoxIndex];
                     selectedBoxInfo.textContent = `#${this.selectedBoxIndex + 1} (${Math.round(box.left)}, ${Math.round(box.top)}, ${Math.round(box.width)}×${Math.round(box.height)})`;
                 } else {
                     selectedBoxInfo.textContent = 'None';
@@ -1280,8 +1629,9 @@
                 };
             }
 
+            // Update the save method to handle new format
             async saveAnnotations() {
-                if (!this.currentImage || this.boxes.length === 0) {
+                if (!this.currentImage || this.annotations.length === 0) {
                     this.showAlert('No annotations to save', 'info');
                     return;
                 }
@@ -1293,10 +1643,7 @@
                             'Content-Type': 'application/json',
                         },
                         body: JSON.stringify({
-                            boxes: this.boxes.map(box => ({
-                                ...box,
-                                saved: true
-                            })),
+                            annotations: this.annotations.map(ann => ({ ...ann, saved: true })),
                             image_name: this.currentImage,
                             original_width: this.originalWidth,
                             original_height: this.originalHeight
@@ -1305,10 +1652,10 @@
 
                     if (response.ok) {
                         const result = await response.json();
-                        this.boxes.forEach(box => box.saved = true);
+                        this.annotations.forEach(box => box.saved = true);
                         this.drawCanvas();
                         this.showAlert(result.message, 'success');
-                        this.loadImages(); // Refresh image list to update progress
+                        this.loadImages();
                     } else {
                         throw new Error('Failed to save annotations');
                     }
@@ -1318,34 +1665,34 @@
             }
 
             undoLastBox() {
-                if (this.boxes.length > 0) {
-                    this.boxes.pop();
+                if (this.annotations.length > 0) {
+                    this.annotations.pop();
                     this.selectedBoxIndex = -1;
-                    document.getElementById('boxCount').textContent = this.boxes.length;
+                    document.getElementById('boxCount').textContent = this.annotations.length;
                     this.updateSelectedBoxInfo();
                     this.drawCanvas();
                     this.showAlert('Last box removed', 'info');
                 } else {
-                    this.showAlert('No boxes to undo', 'info');
+                    this.showAlert('No annotations to undo', 'info');
                 }
             }
 
-            async clearAllBoxes() {
+            async clearAllAnnotations() {
                 if (!this.currentImage) return;
 
-                if (confirm('Are you sure you want to clear all boxes and delete the annotation file?')) {
+                if (confirm('Are you sure you want to clear all annotations and delete the annotation file?')) {
                     try {
                         // Delete annotations from server
                         await fetch(`/api/annotate/annotations/${encodeURIComponent(this.currentImage)}`, {
                             method: 'DELETE'
                         });
 
-                        this.boxes = [];
+                        this.annotations = [];
                         this.selectedBoxIndex = -1;
                         document.getElementById('boxCount').textContent = '0';
                         this.updateSelectedBoxInfo();
                         this.drawCanvas();
-                        this.showAlert('All boxes cleared', 'success');
+                        this.showAlert('All annotations cleared', 'success');
                         this.loadImages(); // Refresh progress
                     } catch (error) {
                         this.showAlert('Error clearing annotations: ' + error.message, 'error');
@@ -1360,12 +1707,12 @@
                     const response = await fetch(`/api/annotate/annotations/${encodeURIComponent(this.currentImage)}`);
                     const data = await response.json();
 
-                    this.boxes = (data.boxes || []).map(box => ({
+                    this.annotations = (data.annotations || []).map(box => ({
                         ...box,
                         saved: true
                     }));
                     this.selectedBoxIndex = -1;
-                    document.getElementById('boxCount').textContent = this.boxes.length;
+                    document.getElementById('boxCount').textContent = this.annotations.length;
                     this.updateSelectedBoxInfo();
                     this.drawCanvas();
                     this.showAlert('Annotations reloaded from file', 'success');
@@ -1381,12 +1728,12 @@
                     const response = await fetch(`/api/annotate/detect_annotations/${encodeURIComponent(this.currentImage)}`);
                     const data = await response.json();
 
-                    this.boxes = (data.boxes || []).map(box => ({
+                    this.annotations = (data.annotations || []).map(box => ({
                         ...box,
                         saved: true
                     }));
                     this.selectedBoxIndex = -1;
-                    document.getElementById('boxCount').textContent = this.boxes.length;
+                    document.getElementById('boxCount').textContent = this.annotations.length;
                     this.updateSelectedBoxInfo();
                     this.drawCanvas();
                     this.showAlert('Annotations detected and loaded from file', 'success');

comic_panel_extractor/utils.py

@@ -7,8 +7,6 @@ import os
 import shutil
 from glob import glob
 from typing import List, Union
-from dotenv import load_dotenv
-load_dotenv()
 from .config import Config
 
 def remove_duplicate_boxes(boxes, compare_single=None, iou_threshold=0.7):

comic_panel_extractor/yolo_manager.py

@@ -3,9 +3,9 @@ import os
 import shutil
 from glob import glob
 from typing import List, Union
-from dotenv import load_dotenv
 
-load_dotenv()
+os.environ["TORCH_USE_CUDA_DSA"] = "1"
+os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
 
 def get_abs_path(relative_path: str) -> str:
     """Convert relative path to absolute path."""