Update app.py

app.py

@@ -1,61 +1,65 @@
 import gradio as gr
-from transformers import AutoImageProcessor, AutoModelForImageClassification
+from transformers import AutoProcessor, AutoModelForCausalLM  # یا AutoModelForVision2Seq اگر vision2seq باشه
 import torch
 from PIL import Image
-import numpy as np
 
-# بارگذاری مدل و processor یک بار موقع شروع
-MODEL_ID = "erfanansaghariyan/mobilew-v11-convnext-tiny-6layer-radiimagen"
+# repo_id دقیق
+MODEL_ID = "erfanasghariyan/mobilew-v11-convnext-tiny-6layer-radimagenet"
 
-processor = AutoImageProcessor.from_pretrained(MODEL_ID)
-model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
+processor = AutoProcessor.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID)   # اگر causalLM نباشه، عوض کن به AutoModelForVision2Seq
 model.eval()
 
-# اگر GPU داشتی فعالش کن
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
-def predict(image: Image.Image):
-    # پیش‌پردازش تصویر
-    inputs = processor(images=image, return_tensors="pt")
-    inputs = {k: v.to(device) for k, v in inputs.items()}
+def generate_report(image: Image.Image, prompt: str = "Describe this radiology image in detail:"):
+    # پردازش تصویر + متن prompt
+    inputs = processor(images=image, text=prompt, return_tensors="pt").to(device)
     
-    # inference
+    # generation
     with torch.no_grad():
-        outputs = model(**inputs)
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=200,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            num_beams=4,  # اگر deterministic می‌خوای beam search
+            repetition_penalty=1.2
+        )
     
-    logits = outputs.logits
-    probabilities = torch.softmax(logits, dim=-1).cpu().numpy()[0]
+    # decode خروجی
+    generated_text = processor.decode(outputs[0], skip_special_tokens=True)
     
-    # گرفتن top-k پیش‌بینی‌ها
-    top_k = 5
-    top_indices = np.argsort(probabilities)[-top_k:][::-1]
-    labels = [model.config.id2label[idx] for idx in top_indices]
-    probs = [float(probabilities[idx]) for idx in top_indices]
+    # اگر prompt در خروجی تکرار شده، تمیز کن
+    if generated_text.startswith(prompt):
+        generated_text = generated_text[len(prompt):].strip()
     
-    # خروجی به صورت دیکشنری برای Gradio
-    result = {label: prob for label, prob in zip(labels, probs)}
-    
-    return result, image  # تصویر + احتمالات
+    return generated_text, image
 
-# رابط کاربری Gradio
 demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil", label="تصویر رادیولوژی آپلود کن"),
+    fn=generate_report,
+    inputs=[
+        gr.Image(type="pil", label="تصویر رادیولوژی آپلود کن (X-ray, CT, MRI و ...)"),
+        gr.Textbox(label="پرامپت (اختیاری)", value="Generate a detailed radiology report for this image:")
+    ],
     outputs=[
-        gr.Label(label="احتمالات کلاس‌ها (Top 5)"),
+        gr.Textbox(label="گزارش / توصیف تولید شده"),
         gr.Image(label="تصویر ورودی")
     ],
-    title="MobileWev-v11 ConvNeXt Tiny - Radiology Image Analysis",
+    title="MobileW-v11 – ConvNeXt Tiny + 6-Layer Decoder for Radiology",
     description=(
-        "این مدل سبک (ConvNeXt-tiny با 6 لایه) برای تحلیل تصاویر پزشکی/رادیولوژی فاین‌تیون شده است.\n"
-        "تصویر را آپلود کنید تا احتمالات کلاس‌های تشخیص را ببینید."
+        "مدل سبک multimodal برای تحلیل تصاویر پزشکی.\n"
+        "encoder: ConvNeXt-Tiny (فریز شده)\n"
+        "decoder: 6 لایه برای تولید متن بهتر\n"
+        "مثال پرامپت‌ها: 'Describe findings', 'Write a radiology report', 'Is there pneumonia?'"
     ),
     examples=[
-        # می‌تونی چند تا مثال تصویر بذاری (آپلود کن توی Space)
-        # مثلاً: ["example_chest_xray.jpg"]
+        # اگر مثال تصویر داری، آپلود کن به Space و مسیر بده
+        # [ "example_chest_xray.jpg", "Write a structured report" ]
     ],
-    cache_examples=False,
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":