Update main.py

main.py

@@ -1,154 +1,229 @@
-from flask import Flask, request, jsonify
-from flask_cors import CORS
-from PIL import Image, ImageDraw
-import io
-import base64
-import time
+import logging
+import aiohttp
+import os,time
+
+from aiohttp import FormData
+from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
+
+from telegram.ext import (
+    Application, CommandHandler, MessageHandler,
+    filters, CallbackQueryHandler, ContextTypes
+)
+
+from dotenv import load_dotenv
+load_dotenv(dotenv_path='env.env')
+
+
+#config
+TOKEN = "8320924107:AAH505mhHkOxeY3aLk0GObIpO_KCtY9hhLM"
+API_ENDPOINT = os.getenv("API_ENDPOINT")
+logging.basicConfig(level=logging.INFO)
+
+#category mappings
+VIEW_CATEGORIES = {
+    "crl": {
+        "Maxilla": "mx",
+        "Mandible-MDS": "mds", 
+        "Mandible-MLS": "mls",
+        "Lateral ventricle": "lv",
+        "Head": "head",
+        "Gestational sac": "gsac",
+        "Thorax": "thorax",
+        "Abdomen": "ab",
+        "Body(Biparietal diameter)": "bd",
+        "Rhombencephalon": "rbp",
+        "Diencephalon": "dp",
+        "NTAPS": "ntaps",
+        "Nasal bone": "nb"
+    },
+    "nt": {
+        "Maxilla": "mx",
+        "Mandible-MDS": "mds", 
+        "Mandible-MLS": "mls",
+        "Lateral ventricle": "lv",
+        "Head": "head",
+        "Thorax": "thorax",
+        "Abdomen": "ab",
+        "Rhombencephalon": "rbp",
+        "Diencephalon": "dp",
+        "Nuchal translucency": "nt",
+        "NTAPS": "ntaps",
+        "Nasal bone": "nb"
+    }
+}
+
+#start
+async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    await update.message.reply_text("Send an ultrasound image (JPG/png only) to begin. See /instructions first")
+
+async def instructions(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    instruction_message = """
+📋 **Instructions**
+
+🔸 **Step 1:** Send a cropped ultrasound image of a structure you want to analyze. See all structures with /list command
+🔸 **Step 2:** Select the ultrasound view (CRL or NT)
+🔸 **Step 3:** Choose the anatomical category to analyze
+🔸 **Step 4:** Wait for the AI analysis results
+
+📌 **Important notes:**
+• Only JPG/PNG images are supported
+• Ensure the ultrasound image is clear and properly oriented
+• Results are for reference only - always consult a medical professional
+• Processing may take a few seconds
+• Stop bot with /stop
+
+💡 **Tips:**
+• Use high-quality, well-lit images for better accuracy
+• Make sure the anatomical structure is clearly visible
+• Different views (CRL/NT) have different category options
+
+🆘 **Need help?** Contact support if you encounter any issues @d3ikshr.
+    """
+    
+    await update.message.reply_text(instruction_message, parse_mode='Markdown')
+
+async def list_categories(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    list_message = """
+📋 **Available Categories by View**
+
+🔍 **CRL view categories:**
+• Maxilla • Mandible-MDS • Mandible-MLS
+• Lateral ventricle • Head • Gestational sac
+• Thorax • Abdomen • Body(Biparietal diameter)
+• Rhombencephalon • Diencephalon • NTAPS
+• Nasal bone
+
+🔍 **NT view categories:**
+• Maxilla • Mandible-MDS • Mandible-MLS
+• Lateral ventricle • Head • Thorax
+• Abdomen • Rhombencephalon • Diencephalon
+• Nuchal translucency • NTAPS • Nasal bone
+
+💡 **Note:** Categories will be shown automatically based on your selected view during analysis.
+    """
+    
+    await update.message.reply_text(list_message, parse_mode='Markdown')
+
+async def stop(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    await update.message.reply_text("🛑 Bot stopped for this chat. Use /start to begin again.")
+    # Clear user data
+    context.user_data.clear()
+
+# Receive image
+async def handle_image(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    photo = update.message.photo[-1]
+    file = await photo.get_file()
+    file_path = f"{update.message.from_user.id}_ultrasound.jpg"
+    await file.download_to_drive(file_path)
+    context.user_data["image_path"] = file_path
+    
+    # get view
+    buttons = [
+        [InlineKeyboardButton("CRL", callback_data="view:crl"),
+         InlineKeyboardButton("NT", callback_data="view:nt")]
+    ]
+    await update.message.reply_text(
+        "Select the ultrasound view:",
+        reply_markup=InlineKeyboardMarkup(buttons)
+    )
+
+#selcet view
+async def handle_view(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    query = update.callback_query
+    await query.answer()
+    view = query.data.split(":")[1]
+    context.user_data["selected_view"] = view
+    
+    #get categories for the selected view
+    categories = list(VIEW_CATEGORIES[view].keys())
+    buttons = [[InlineKeyboardButton(cat, callback_data=f"category:{cat}")]
+               for cat in categories]
+    await query.edit_message_text(
+        "Select anatomical category:",
+        reply_markup=InlineKeyboardMarkup(buttons)
+    )
+
+#get category then upload
+async def handle_category(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    query = update.callback_query
+    await query.answer()
+    category_display = query.data.split(":")[1]
+    context.user_data["selected_category"] = category_display
+    
+    image_path = context.user_data.get("image_path")
+    view = context.user_data.get("selected_view")
+    
+    if not image_path or not view:
+        await query.edit_message_text("Missing image or view.")
+        return
+
+    await query.edit_message_text("🔄 Processing image...")
+    time.sleep(3)
+    await query.edit_message_text("🏥 Building diagnosis, please wait...")
 
-from model import predict
-
-app = Flask(__name__)
-CORS(app)
-
-# global variables
-threshold = 0.02
-diagnosis = ''
-pc_source = False
-
-@app.route('/api/process-annotation', methods=['POST'])
-def process_annotation():
     try:
-        if request.content_type and request.content_type.startswith('multipart/form-data'):
-            source = request.form.get('source', 'telegram')
-        else:
-            data = request.get_json()
-            source = data.get('source', 'pc') if data else 'pc'
-
-        if source == "pc":
-            global pc_source
-            pc_source = True
-
-            if not data:
-                return jsonify({"error": "No JSON data provided"}), 400
-            img_data = data.get('image')
-            bbox = data.get('coordinates')
-            category = data.get('category')
-            view = data.get('view')
-
-            header, encoded = img_data.split(',', 1)
-            img_bytes = base64.b64decode(encoded)
-            img = Image.open(io.BytesIO(img_bytes))
-
-            cropped = img.crop((
-                bbox['x'],
-                bbox['y'],
-                bbox['x'] + bbox['width'],
-                bbox['y'] + bbox['height']
-            ))
-
-            print(f"\nParameters:[Category-{category}, View-{view}, Annotations-{bbox}]\n")
-            print("Processing...\n")
-            time.sleep(2)
-
-            error = predict(cropped, view, category)
-            label = get_label(error)
-            confidence = 100 - (100 * float(error))
-            diagnosis = build_diagnosis(error, confidence)
-
-            draw = ImageDraw.Draw(img)
-            draw.rectangle(
-                [bbox['x'], bbox['y'], bbox['x'] + bbox['width'], bbox['y'] + bbox['height']],
-                outline="red", width=3
-            )
-
-            buffered = io.BytesIO()
-            cropped.save(buffered, format="PNG")
-            processed_b64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
-            processed_data_url = f"data:image/png;base64,{processed_b64}"
-
-            print(f"Results: Error-{error}, Comment-{label}, Confidence-{confidence}, Diagnosis-{diagnosis}\n")
-            print(f"Rendering results to: {source} \n")
-
-            return jsonify({
-                "processed_image": processed_data_url,
-                "category": category,
-                "comment": label,
-                "error": error,
-                "confidence": confidence,
-                "threshold": threshold,
-                "diagnosis": diagnosis
-            })
-
-        elif source == "telegram":
-            category = request.form.get('category')
-            view = request.form.get('view')
-            source_field = request.form.get('source')
-            image_file = request.files.get('image')
-
-            if not image_file:
-                return jsonify({"error": "No image file provided"}), 400
-
-            if not category or not view:
-                return jsonify({"error": "Missing category or view"}), 400
-
-            try:
-                img = Image.open(image_file)
-                if img.mode != 'RGB':
-                    img = img.convert('RGB')
-            except Exception as e:
-                return jsonify({"error": f"Invalid image file: {str(e)}"}), 400
-
-            print(f"\nParameters:[Category-{category}, View-{view}, Image size-{img.size}]\n")
-            print("Processing...\n")
-
-            error = predict(img, view, category)
-            label = get_label(error)
-            confidence = 100 - (100 * float(error))
-            diagnosis = build_diagnosis(error, confidence)
-
-            buffered = io.BytesIO()
-            img.save(buffered, format="PNG")
-            processed_b64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
-            processed_data_url = f"data:image/png;base64,{processed_b64}"
-
-            response_data = {
-                "processed_image": processed_data_url,
-                "category": category,
-                "view": view,
-                "comment": label,
-                "error": error,
-                "confidence": confidence,
-                "threshold": threshold,
-                "diagnosis": diagnosis,
-                "source": "telegram"
-            }
-
-            print(f"Results: Error-{error}, Comment-{label}, Confidence-{confidence}, Diagnosis-{diagnosis}\n")
-            print(f"Rendering results to: {source} \n")
-
-            return jsonify(response_data)
-
-        else:
-            return jsonify({"error": "Unsupported source"}), 400
-
+        #read the image file into memory first
+        with open(image_path, "rb") as f:
+            image_data = f.read()
+        
+        #mapping
+        category_value = VIEW_CATEGORIES[view][category_display]
+        
+        #create form data
+        form = FormData()
+        form.add_field("view", view)
+        form.add_field("category", category_value)
+        form.add_field("source", "telegram")
+        form.add_field("image", image_data, filename="image.jpg", content_type="image/jpeg")
+        
+        #send request 
+        async with aiohttp.ClientSession() as session:
+            async with session.post(API_ENDPOINT, data=form) as resp:
+                if resp.status == 200:
+                    result = await resp.json()
+                    
+                    message = f"""
+📊 **Analysis Results**
+
+🔍 **View:** {result.get('view', view).upper()}
+🏥 **Category:** {category_display}
+
+📈 **Confidence:** {result.get('confidence', 0):.2f}%
+⚠️ **Reconstruction error:** {result.get('error', 0):.5f}
+
+📋 **Status:** {result.get('comment', 'No comment')}
+
+🩺 **Diagnosis:** {result.get('diagnosis', 'No diagnosis')}*
+                    """
+                    
+                    await query.edit_message_text(message, parse_mode='Markdown')
+                     
+                else:
+                    error_text = await resp.text()
+                    await query.edit_message_text(f"❌ Upload failed. Status: {resp.status}\nError: {error_text}")
+                    
     except Exception as e:
-        print(f"Unexpected error: {e}")
-        return jsonify({"error": str(e)}), 500
-
-def get_label(error):
-    return "⚠️ Anomaly detected" if error > threshold else "✅ Normal structure"
-
-def build_diagnosis(error, confidence):
-    if confidence < 100 and confidence > 98 and error > threshold:
-        diagnosis = 'High probability that there is an anomaly in the structure.'
-    elif confidence < 98 and confidence > 96 and error > threshold:
-        diagnosis = 'My analysis concludes there could be an anomaly in the structure, but either the structure is not well annotated or the image could be distorted, small, unclear or bad hence the uncertainty.'
-    elif confidence < 96 and confidence > 92 and error > threshold:
-        diagnosis = 'High possibility that anomaly is false positive and image is greatly distorted or irrelevant. Please check the image or the annotation and try again. If results are same, consult an expert.'
-    elif confidence < 92 and error > threshold:
-        diagnosis = 'Please upload a good ultrasound scan to obtain diagnosis. I cannot recognize the image nor the outlined structure.'
-    elif confidence < 100 and confidence > 98 and error < threshold:
-        diagnosis = 'Healthy structure detected. Annotation is correct OR model partially detects healthy area.'
-
-    diagnosis += " THIS IS NOT PROFESSIONAL MEDICAL ADVICE. LIASE WITH AN EXPERT"
-    return diagnosis
+        logging.error(f"Error processing request: {e}")
+        await query.edit_message_text(f"❌ Error: {str(e)}")
+        
+    finally:
+        try:
+            if os.path.exists(image_path):
+                os.remove(image_path)
+                logging.info(f"Cleaned up image file: {image_path}")
+        except Exception as e:
+            logging.error(f"Error cleaning up file: {e}")
+
+def main():
+    app = Application.builder().token(TOKEN).build()
+    app.add_handler(CommandHandler("start", start))
+    app.add_handler(CommandHandler("instructions", instructions))
+    app.add_handler(CommandHandler("list", list_categories))
+    app.add_handler(CommandHandler("stop", stop))
+    app.add_handler(MessageHandler(filters.PHOTO, handle_image))
+    app.add_handler(CallbackQueryHandler(handle_view, pattern="^view:"))
+    app.add_handler(CallbackQueryHandler(handle_category, pattern="^category:"))
+    app.run_polling()
+
+if __name__ == "__main__":
+    main()