fix

README.md

@@ -1,6 +1,6 @@
 ---
-title: Phronesis
-emoji: 🌖
+title: Phronesis Medical Report Generator
+emoji: 🧠
 colorFrom: green
 colorTo: gray
 sdk: gradio
@@ -8,6 +8,82 @@ sdk_version: 5.4.0
 app_file: app.py
 pinned: false
 short_description: 'REPORT GEN AND CLASSIFICATION MODEL '
+
+---
+
+# 🧠 Phronesis: Medical Image Diagnosis & Report Generator
+
+**Phronesis** is a multimodal AI tool that classifies medical CT scan images (DICOM or standard formats) and generates diagnostic reports using a combination of video classification and medical language generation.
+
+---
+
+## 🚀 Demo
+
+Upload a set of DICOM (`.dcm`, `.ima`) or image (`.png`, `.jpg`) files representing slices of a CT scan. The model will:
+
+- 🏷️ Predict a class: **acute**, **normal**, **chronic**, or **lacunar**
+- 📋 Generate a short **radiology report**
+
+[Live App →](https://huggingface.co/spaces/baliddeki/phronesis-ml-endpoint)
+
+---
+
+## 🏗️ Model Architecture
+
+- **Vision Backbone**: `3D ResNet-18` pretrained on Kinetics-400
+- **Language Head**: `BioBART v2` (pretrained biomedical seq2seq model)
+- **Bridge Module**: Custom `ImageToTextProjector` to align visual features with the language model
+- **CombinedModel**: Unified architecture for classification + report generation
+
+---
+
+## 🧪 Tasks
+
+- **Image Classification**: Categorizes brain CT scans into one of four classes.
+- **Report Generation**: Produces diagnostic text conditioned on image features.
+
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## 🖼️ Input Format
+
+- Minimum 1, maximum ~30 image slices per scan.
+- Acceptable file formats:
+  - DICOM (`.dcm`, `.ima`)
+  - PNG, JPEG
+
+The model will sample or pad the series to 16 frames for temporal context.
+
+---
+
+## 📦 Dependencies
+
+This app uses:
+- `torch`
+- `transformers`
+- `torchvision`
+- `huggingface_hub`
+- `pydicom`
+- `gradio`
+- `PIL`, `numpy`
+
+---
+
+## 🔐 Notes
+
+- This demo loads a private model from the Hugging Face Hub. Set your `HF_TOKEN` as a secret for the space if needed.
+- Do **not use for real clinical decisions** – intended for research/demo only.
+
+---
+
+## 🙋‍♂️ Credits
+
+Developed by [@baliddeki](https://huggingface.co/baliddeki)
+
+Model weights: [`baliddeki/phronesis-ml`](https://huggingface.co/baliddeki/phronesis-ml)  
+Language model: [`GanjinZero/biobart-v2-base`](https://huggingface.co/GanjinZero/biobart-v2-base)
+
+---
+
+## 📄 License
+
+MIT or Apache 2.0 (add yours here)

app.py

@@ -1,129 +1,114 @@
-#app.py
-import os
-import io
-import uvicorn
-
+# app.py
 import torch
-from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.responses import JSONResponse
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from torchvision import models, transforms
-from PIL import Image
 import numpy as np
+from PIL import Image
+import io
+import gradio as gr
+from torchvision import models, transforms
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from huggingface_hub import hf_hub_download
+from model import CombinedModel, ImageToTextProjector
 import pydicom
+import os
 import gc
-from model import CombinedModel, ImageToTextProjector
-from fastapi import FastAPI, Request
-from fastapi.middleware.cors import CORSMiddleware
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-app = FastAPI()
+# Load tokenizer and models
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+tokenizer = AutoTokenizer.from_pretrained("baliddeki/phronesis-ml", token=HF_TOKEN)
+video_model = models.video.r3d_18(weights="KINETICS400_V1")
+video_model.fc = torch.nn.Linear(video_model.fc.in_features, 512)
 
-@app.get("/")
-async def root(request: Request):
-    return {"message": "Welcome to Phronesis"}
+report_generator = AutoModelForSeq2SeqLM.from_pretrained("GanjinZero/biobart-v2-base")
+projector = ImageToTextProjector(512, report_generator.config.d_model)
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+num_classes = 4
+class_names = ["acute", "normal", "chronic", "lacunar"]
+combined_model = CombinedModel(
+    video_model, report_generator, num_classes, projector, tokenizer
+)
+model_file = hf_hub_download(
+    "baliddeki/phronesis-ml", "pytorch_model.bin", token=HF_TOKEN
+)
+state_dict = torch.load(model_file, map_location=device)
+combined_model.load_state_dict(state_dict)
+combined_model.eval()
+
+# Image transforms
+image_transform = transforms.Compose(
+    [
+        transforms.Resize((112, 112)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.43216, 0.394666, 0.37645], std=[0.22803, 0.22145, 0.216989]
+        ),
+    ]
+)
 
-def dicom_to_png(dicom_data):
-    try:
-        dicom_file = pydicom.dcmread(dicom_data)
-        if not hasattr(dicom_file, 'PixelData'):
-            raise HTTPException(status_code=400, detail="No pixel data in DICOM file.")
-        
-        pixel_array = dicom_file.pixel_array.astype(np.float32)
-        pixel_array = ((pixel_array - pixel_array.min()) / (pixel_array.ptp())) * 255.0
-        pixel_array = pixel_array.astype(np.uint8)
-
-        img = Image.fromarray(pixel_array).convert("L")
-        return img
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error converting DICOM to PNG: {e}")
-
-# Set up secure model initialization
-HF_TOKEN = os.getenv('HF_TOKEN')
-if not HF_TOKEN:
-    raise ValueError("Missing Hugging Face token in environment variables.")
-
-try:
-    report_generator_tokenizer = AutoTokenizer.from_pretrained(
-        "baliddeki/phronesis-ml",
-        token=HF_TOKEN if HF_TOKEN else None
-    )
-    video_model = models.video.r3d_18(weights="KINETICS400_V1")
-    video_model.fc = torch.nn.Linear(video_model.fc.in_features, 512)
-    report_generator = AutoModelForSeq2SeqLM.from_pretrained("GanjinZero/biobart-v2-base")
-    projector = ImageToTextProjector(512, report_generator.config.d_model)
-    num_classes = 4
-    combined_model = CombinedModel(video_model, report_generator, num_classes, projector, report_generator_tokenizer)
-    model_file = hf_hub_download("baliddeki/phronesis-ml", "pytorch_model.bin", token=HF_TOKEN)
-    state_dict = torch.load(model_file, map_location=device)
-    combined_model.load_state_dict(state_dict)
-    combined_model.eval()
-except Exception as e:
-    raise SystemExit(f"Error loading models: {e}")
-
-image_transform = transforms.Compose([
-    transforms.Resize((112, 112)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.43216, 0.394666, 0.37645], std=[0.22803, 0.22145, 0.216989])
-])
 
-class_names = ["acute", "normal", "chronic", "lacunar"]
+def dicom_to_image(file_bytes):
+    dicom_file = pydicom.dcmread(io.BytesIO(file_bytes))
+    pixel_array = dicom_file.pixel_array.astype(np.float32)
+    pixel_array = ((pixel_array - pixel_array.min()) / pixel_array.ptp()) * 255.0
+    pixel_array = pixel_array.astype(np.uint8)
+    return Image.fromarray(pixel_array).convert("RGB")
 
-@app.post("/predict/")
-async def predict(files: list[UploadFile]):
-    print(f"Received {len(files)} files")
-    n_frames = 16
-    images = []
-
-    for file in files:
-        ext = file.filename.split('.')[-1].lower()
-        try:
-            if ext in ['dcm', 'ima']:
-                dicom_img = dicom_to_png(await file.read())
-                images.append(dicom_img.convert("RGB"))
-            elif ext in ['png', 'jpeg', 'jpg']:
-                img = Image.open(io.BytesIO(await file.read())).convert("RGB")
-                images.append(img)
-            else:
-                raise HTTPException(status_code=400, detail="Unsupported file type.")
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error processing file {file.filename}: {e}")
 
+def predict(images):
     if not images:
-        return JSONResponse(content={"error": "No valid images provided."}, status_code=400)
-
-    if len(images) >= n_frames:
-        images_sampled = [images[i] for i in np.linspace(0, len(images) - 1, n_frames, dtype=int)]
+        return "No image uploaded.", ""
+
+    # Convert images
+    processed_imgs = []
+    for img in images:
+        filename = img.name.lower()
+        if filename.endswith((".dcm", ".ima")):
+            dicom_img = dicom_to_image(img.read())
+            processed_imgs.append(dicom_img)
+        else:
+            pil_img = Image.open(img).convert("RGB")
+            processed_imgs.append(pil_img)
+
+    # Sample or pad
+    n_frames = 16
+    if len(processed_imgs) >= n_frames:
+        images_sampled = [
+            processed_imgs[i]
+            for i in np.linspace(0, len(processed_imgs) - 1, n_frames, dtype=int)
+        ]
     else:
-        images_sampled = images + [images[-1]] * (n_frames - len(images))
+        images_sampled = processed_imgs + [processed_imgs[-1]] * (
+            n_frames - len(processed_imgs)
+        )
 
-    image_tensors = [image_transform(img) for img in images_sampled]
-    images_tensor = torch.stack(image_tensors).permute(1, 0, 2, 3).unsqueeze(0).to(device)
+    tensor_imgs = [image_transform(i) for i in images_sampled]
+    input_tensor = torch.stack(tensor_imgs).permute(1, 0, 2, 3).unsqueeze(0).to(device)
 
     with torch.no_grad():
-        class_outputs, generated_report, _ = combined_model(images_tensor)
-        predicted_class = torch.argmax(class_outputs, dim=1).item()
-        predicted_class_name = class_names[predicted_class]
+        class_logits, report, _ = combined_model(input_tensor)
+        class_pred = torch.argmax(class_logits, dim=1).item()
+        class_name = class_names[class_pred]
 
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-        
-    return {
-        "predicted_class": predicted_class_name,
-        "generated_report": generated_report[0] if generated_report else "No report generated."
-    }
-
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))
+
+    return class_name, report[0] if report else "No report generated."
+
+
+# Gradio interface
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.File(
+        file_types=[".dcm", ".jpg", ".jpeg", ".png"],
+        file_count="multiple",
+        label="Upload CT Scan Images",
+    ),
+    outputs=[gr.Textbox(label="Predicted Class"), gr.Textbox(label="Generated Report")],
+    title="Phronesis Medical Report Generator",
+    description="Upload CT scan DICOM or image files. Returns diagnosis classification and generated report.",
+)
+
+demo.launch()