app.py

from fastapi import FastAPI, UploadFile, status
from fastapi.responses import JSONResponse
from ultralytics import YOLO
import cv2
import numpy as np
from fastapi.middleware.cors import CORSMiddleware
import base64

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["POST"],
    allow_headers=["content-type", "accept"],
)

model = YOLO("best.pt")

@app.post("/inference")
async def inference(file: UploadFile):
    if file.content_type != "image/jpeg" and file.content_type != "image/png" and file.content_type != "image/jpg":
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={"error": "Invalid file format"}
        )
    image_bytes = await file.read()
    img = np.frombuffer(image_bytes, dtype=np.uint8)
    img = cv2.imdecode(img, cv2.IMREAD_COLOR)
    results = model.predict(source=img, conf=0.3)
    detections = []
    for r in results: 
        boxes = r.boxes
    for box in boxes: 
        x1, y1, x2, y2 = box.xyxy[0]
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
    
        cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 255), 2)

        label = box.cls[0].item()  
        label_name = model.names[label]  
        detections.append(label_name)
        (text_width, text_height), baseline = cv2.getTextSize(label_name, cv2.FONT_HERSHEY_SIMPLEX, 0.9, 2)
        cv2.rectangle(img, (x1, y1 - text_height - baseline - 5), (x1 + text_width, y1), (255, 0, 255), -1)
        cv2.putText(img, label_name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 0), 2)
    
    resp_img_bytes = cv2.imencode('.jpg', img)[1].tobytes()
    img_base64 = base64.b64encode(resp_img_bytes).decode('utf-8')
    return JSONResponse({
        "image": img_base64,
        "detections": detections
    })


AVERAGE_TUMOR_VOLUME = 523.6 

def calculate_sphere_volume(width, height):
    """
    Menghitung volume tumor menggunakan rumus bola
    Volume = (4/3) * π * r³
    Diameter = rata-rata dari width dan height
    """
    try:
        # Diameter sebagai rata-rata dari width dan height (dalam pixel)
        diameter = (width + height) / 2
        radius = diameter / 2
        
        # Rumus volume bola
        volume = (4/3) * math.pi * (radius ** 3)
        
        return round(volume, 2)
    except:
        return None

@app.post("/inference_volume")
async def inference_volume(file: UploadFile):
    """
    Endpoint untuk deteksi tumor dengan volume, return image dengan anotasi
    """
    if file.content_type not in ["image/jpeg", "image/png", "image/jpg"]:
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={"error": "Invalid file format"}
        )
    
    image_bytes = await file.read()
    img = np.frombuffer(image_bytes, dtype=np.uint8)
    img = cv2.imdecode(img, cv2.IMREAD_COLOR)
    
    results = model.predict(
        source=img, 
        conf=0.5,
        iou=0.2
    )
    
    total_volume = 0
    detection_count = 0
    
    for r in results: 
        boxes = r.boxes
        
        for box in boxes: 
            x1, y1, x2, y2 = box.xyxy[0]
            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
            
            width = x2 - x1
            height = y2 - y1
            
            volume = calculate_sphere_volume(width, height)
            if volume is None:
                volume = AVERAGE_TUMOR_VOLUME
            
            total_volume += volume
            detection_count += 1
            
            label = int(box.cls[0].item())
            label_name = model.names[label]
            confidence = box.conf[0].item()
            
            # Draw bounding box
            cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 255), 2)
            
            # Label dengan volume
            text = f"{label_name} {confidence:.2f}"
            vol_text = f"Vol: {volume:.1f}mm3"
            
            (text_width, text_height), baseline = cv2.getTextSize(
                text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2
            )
            
            cv2.rectangle(
                img, 
                (x1, y1 - text_height - baseline - 25), 
                (x1 + max(text_width, 120), y1), 
                (255, 0, 255), 
                -1
            )
            
            cv2.putText(
                img, text, (x1, y1 - 20), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2
            )
            cv2.putText(
                img, vol_text, (x1, y1 - 5), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1
            )
    
    summary_text = f"Total: {detection_count} tumor(s) | Vol: {total_volume:.1f}mm3"
    cv2.rectangle(img, (10, 10), (400, 40), (0, 0, 0), -1)
    cv2.putText(img, summary_text, (15, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
    
    resp_img_bytes = cv2.imencode('.jpg', img)[1].tobytes()
    resp_filename = f"volume_{file.filename}" if file.filename else "volume_image.jpg"
    
    return StreamingResponse(
        BytesIO(resp_img_bytes), 
        media_type="image/jpeg", 
        headers={"Content-Disposition": f"attachment; filename={resp_filename}"}
    )