rafactor code for images

app.py

@@ -108,16 +108,7 @@ async def start_detection(data: Dict):
     model_path = model_by_id_metadata(parts['id'])
 
     logger.info(f"[INFO] Checking model_path")
-    if isinstance(model_path, str):
-        if not os.path.exists(model_path):
-            logger.error(f"[ERROR] Model file not found: {model_path}")
-            return {
-                "status": "error",
-                "message": f"Model file not found: {model_path}"
-            }
-        model = YOLO(model_path)
-    else:
-        model = model_path
+    model = model_path
 
     # =====================================================
     # BASE64 IMAGE DETECTION (NOT VIDEO STREAM)

core/defect_detection_image.py

@@ -1,5 +1,8 @@
-import os, cv2, time, base64, asyncio, httpx
+import os, cv2, base64, asyncio, httpx
 import numpy as np
+from ultralytics import YOLO
+from PIL import Image
+from io import BytesIO
 
 from datetime import datetime
 from dotenv import load_dotenv
@@ -16,18 +19,25 @@ WEBHOOK_TIMEOUT = float(os.getenv("WEBHOOK_TIMEOUT", "10.0"))
 # ============================================================
 # DEFECT DETECTION FROM BASE64 IMAGE
 # ============================================================
-def detect_defect_from_base64(station_id: str, camera_id: str, image_base64: str, model=None):
+def detect_defect_from_base64(station_id: str, camera_id: str, image_base64: str, model_path=None):
     """
     Detect defect from a single Base64 image.
-    - Decode Base64 → OpenCV image
-    - Run YOLO model once
-    - Return result (OK / NG) + annotated image base64
+    Return:
+        - status: "OK" / "NG" / "error"
+        - annotated image (base64)
+        - list of detections
     """
-
     try:
-        img_data = base64.b64decode(image_base64)
-        np_arr = np.frombuffer(img_data, np.uint8)
-        frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+        # OPTION 1
+        # img_data = base64.b64decode(image_base64)
+        # np_arr = np.frombuffer(img_data, np.uint8)
+        # frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+        
+        # OPTION 2
+        img_data = base64.b64decode(image_base64) 
+        image = Image.open(BytesIO(img_data)).convert("RGB") 
+        frame = np.array(image)
+        
         if frame is None:
             raise ValueError("Decoded image is None")
     except Exception as e:
@@ -39,128 +49,114 @@ def detect_defect_from_base64(station_id: str, camera_id: str, image_base64: str
             "status_defect": "",
             "image_base64": "",
             "detections": [],
-            "message": f"Invalid base64 image"
+            "message": "Invalid base64 image"
         }
-        
-        
-    # ========== YOLO DETECTION ==========
+
+    detections = []
+
+    try:
+        model = YOLO(f"./{model_path}") 
+        logger.info(f"[MODEL] Success load model")
+    except Exception as e:
+        logger.error(f"[ERROR] Cannot load model: {e}")
+
     if model:
-        # Predict
         results = model.predict(source=frame, conf=0.4, imgsz=640, verbose=False)
         boxes = results[0].boxes
 
-        if len(boxes) > 0:
-            for box in boxes:
-                cls = int(box.cls[0])
-                conf = float(box.conf[0])
-                xyxy = [int(x) for x in box.xyxy[0].tolist()]
-                defect_name = model.names.get(cls, f"class_{cls}").lower()
-
-                x1, y1, x2, y2 = xyxy
-
-                # Ambil warna berdasarkan defect
-                try:
-                    color = color_defect(defect_name)
-                except Exception:
-                    color = color_defect('other')
-
-                # Draw bounding box di frame
-                cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
-
-                # Label
-                label = f"{defect_name.upper()} {conf:.2f}"
-                (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
-                cv2.rectangle(frame, (x1, y1 - 20), (x1 + w, y1), color, -1)
-                cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
-
-
-                # Convert annotated image ke Base64
-                _, buffer = cv2.imencode(".jpg", frame)
-                frame_base64 = base64.b64encode(buffer).decode("utf-8")
-
-                # Save annotated image
-                # output_dir = "outputs/images"
-                # os.makedirs(output_dir, exist_ok=True)
-                # filename = f"{station_id}_{camera_id}_NG_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
-                # filepath = os.path.join(output_dir, filename)
-                # cv2.imwrite(filepath, frame)
-
-                # logger.info(f"[SAVED] NG image saved to {filepath}")
-
-                logger.info(f"[DETECTED] Camera {camera_id} → {defect_name} ({conf:.2f})")
-
-                return {
-                    "station_id": station_id,
-                    "camera_id": camera_id,
-                    "status": "success",
-                    "status_defect": "NG",
-                    "image_base64": frame_base64,
-                    "detections": [{
-                        "class": defect_name,
-                        "confidence": conf,
-                        "bbox": xyxy
-                    }],
-                    "message": f"Detected as defect"
-                }
-
-    # ========== NO DEFECT ==========
+        for box in boxes:
+            cls = int(box.cls[0])
+            conf = float(box.conf[0])
+            xyxy = [int(x) for x in box.xyxy[0].tolist()]
+            defect_name = model.names.get(cls, f"class_{cls}").lower()
+
+            x1, y1, x2, y2 = xyxy
+            color = color_defect(defect_name) if defect_name else color_defect('other')
+
+            # Draw bbox + label
+            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
+            label = f"{defect_name.upper()} {conf:.2f}"
+            (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
+            cv2.rectangle(frame, (x1, y1 - 20), (x1 + w, y1), color, -1)
+            cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 2)
+
+            detections.append({
+                "class": defect_name,
+                "confidence": conf,
+                "bbox": xyxy
+            })
+            
+    # Convert annotated frame ke Base64
     _, buffer = cv2.imencode(".jpg", frame)
     frame_base64 = base64.b64encode(buffer).decode("utf-8")
 
     # Save OK image (no bbox)
-    # output_dir = "outputs/images"
-    # os.makedirs(output_dir, exist_ok=True)
-    # filename = f"{station_id}_{camera_id}_OK_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
-    # filepath = os.path.join(output_dir, filename)
-    # cv2.imwrite(filepath, frame)
-    # logger.info(f"[SAVED] OK image saved to {filepath}")
-
-    logger.info(f"[OK] Camera {camera_id} → No defect detected.")
-    return {
-        "station_id": station_id,
-        "camera_id": camera_id,
-        "status": "success",
-        "status_defect": "OK",
-        "image_base64": frame_base64,
-        "detections": [],
-        "message": f"Detected as normal (no defect)"
-    }
+    output_dir = "outputs/images"
+    os.makedirs(output_dir, exist_ok=True)
+    filename = f"{station_id}_{camera_id}_OK_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
+    filepath = os.path.join(output_dir, filename)
+    cv2.imwrite(filepath, frame)
+    logger.info(f"[SAVED] OK image saved to {filepath}")
+
+    if detections:
+        logger.info(f"[DETECTED] Camera {camera_id} → {len(detections)} defect(s)")
+        return {
+            "station_id": station_id,
+            "camera_id": camera_id,
+            "status": "success",
+            "status_defect": "NG",
+            "image_base64": frame_base64,
+            "detections": detections,
+            "message": "Detected as defect"
+        }
+    else:
+        logger.info(f"[OK] Camera {camera_id} → No defect detected.")
+        return {
+            "station_id": station_id,
+            "camera_id": camera_id,
+            "status": "success",
+            "status_defect": "OK",
+            "image_base64": frame_base64,
+            "detections": [],
+            "message": "Detected as normal (no defect)"
+        }
 
 # ============================================================
 # ASYNC WRAPPERS
 # ============================================================
-async def _detect_camera_image(station_id: str, camera: Dict, model=None):
+async def _detect_camera_image(station_id: str, camera: Dict, model_path=None):
     """Run detect_defect_from_base64 in thread for async parallel."""
     return await asyncio.to_thread(
         detect_defect_from_base64,
         station_id,
         camera["camera_id"],
         camera["image_base64"],
-        model
+        model_path
     )
+    # return await asyncio.to_thread(
+    #     testing,
+    #     station_id,
+    #     camera["camera_id"],
+    #     camera["image_base64"],
+    #     model
+    # )
 
 async def run_detection_group(
     station_id: str,
     cameras: List[Dict],
     webhook_url: str,
-    model=None,
+    model_path=None,
     parts: Dict = None
 ):
-    """
-    Run defect detection for multiple cameras in parallel and send webhook.
-    All results are serialized safely for JSON.
-    """
     parts = parts or {}
-
     logger.info(f"[START] Station {station_id} → {len(cameras)} camera(s)")
 
-    # Run detection async parallel
     results = await asyncio.gather(
-        *[_detect_camera_image(station_id, cam, model) for cam in cameras],
+        *[_detect_camera_image(station_id, cam, model_path) for cam in cameras],
         return_exceptions=True
     )
 
-    # Convert exceptions to dict
+    # Bersihkan exception
     clean_results = []
     for r in results:
         if isinstance(r, Exception):
@@ -171,31 +167,25 @@ async def run_detection_group(
         else:
             clean_results.append(r)
 
-    # Determine overall status
     has_error = any(r.get("status") == "error" for r in clean_results)
     all_error = all(r.get("status") == "error" for r in clean_results)
 
+    status = "success"
+    message = "Success detecting defects"
     if all_error:
-        status = "error"
-        message = "All cameras failed during detection"
+        status, message = "error", "All cameras failed during detection"
     elif has_error:
-        status = "partial_error"
-        message = "Some cameras failed during detection"
-    else:
-        status = "success"
-        message = "Success detecting defects"
+        status, message = "partial_error", "Some cameras failed during detection"
 
-    # Build payload
     payload = {
         "status": status,
-        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime()),
+        "timestamp": datetime.now().isoformat(),
         "model_version": MODEL_VERSION,
         "message": message,
         "parts": parts,
-        # "data": make_serializable(clean_results),
+        "data": make_serializable(clean_results),
     }
 
-    # Send webhook
     try:
         async with httpx.AsyncClient(timeout=WEBHOOK_TIMEOUT) as client:
             await client.post(webhook_url, json=payload)