Update main.py

main.py

@@ -1,276 +1,84 @@
-from enkacard import encbanner
-import asyncio
 import os
-import concurrent.futures
-
-# from enkanetwork import EnkaNetworkAPI
-from fastapi import FastAPI
-from io import BytesIO
-import enkanetwork
-from fastapi.responses import (
-    JSONResponse,
-)  
-import enkacard
-import uvicorn
-import requests
-
-# client = EnkaNetworkAPI()
-
-# import asyncio
-
-# async def card():
-#     print("cret")
-#     async with encbanner.ENC(uid = "811455610") as encard:
-#         print("creting")
-#         return await encard.creat()
-
-# result = asyncio.run(card()) 
-
-# print(result)
-
-# data_dir = "/tmp/data"
-# if not os.path.exists(data_dir):
-#     os.makedirs(data_dir)
-# data_dir = "/tmp/langs"
-# if not os.path.exists(data_dir):
-#     os.makedirs(data_dir)
-
-# async def main():
-#     async with client:
-#         await client.update_assets(lang = ["EN", "CHT"], path="/tmp")
-
-# asyncio.run(main())
-
-app = FastAPI()
-# app = FastAPI(lifespan=lifespan)
-async def card(id, designtype):
-    async with encbanner.ENC(uid = str(id)) as encard:
-        return await encard.creat(template = (2 if str(designtype) == "2" else 1))
+import uuid
+import tempfile
+import httpx
+import torch
+from PIL import Image
+from fastapi import FastAPI, Query, HTTPException
+from transformers import AutoModelForImageClassification, ViTImageProcessor
+from typing import Optional
+
+# Initialize the model and processor globally to avoid reloading for each request
+model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
+processor = ViTImageProcessor.from_pretrained('Falconsai/nsfw_image_detection')
+
+app = FastAPI(title="NSFW Image Detection API")
+
+@app.get("/detect")
+async def detect_nsfw(
+    url: str = Query(..., description="URL of the image to analyze"),
+    timeout: Optional[int] = Query(default=10, description="Timeout for downloading image in seconds")
+):
+    """
+    Detect NSFW content in an image from a given URL.
     
-@app.get("/{id}")  # Correct route definition without prefix
-async def characters(id: int, design: str = "1"):  # Use async and await
+    - Saves the image to a temporary file
+    - Processes the image using a pre-trained NSFW detection model
+    - Returns the highest confidence prediction
+    """
     try:
-        # characters = []
-        result = await card(id, design)  # Use await instead of asyncio.run()
-        print(result)
-        characters = process_images(result)
-        print(characters)
-
-        # Return valid JSON response using FastAPI's JSONResponse
-        return JSONResponse(content={'response': characters})
-
-    except enkanetwork.exception.VaildateUIDError:
-        return JSONResponse(content={'error': 'Invalid UID. Please check your UID.'}, status_code=400)
-
-    except enkacard.enc_error.ENCardError:
-        return JSONResponse(content={'error': 'Enable display of the showcase in the game or add characters there.'}, status_code=400)
-
-    except Exception as e:
-        return JSONResponse(content={'error': 'UNKNOWN ERR: ' + str(e)}, status_code=500)
-
-@app.get("/")
-def hello_world():
-    return 'AMERICA YA HALLO!!'
-
-# @app.route("/update_char")
-# def upload():
-    # data_dir = "/tmp/data"
-
-    # if not os.path.exists(data_dir):
-    #     os.makedirs(data_dir)
-
-    # data_dir = "/tmp/langs"
-
-    # if not os.path.exists(data_dir):
-    #     os.makedirs(data_dir)
-    # asyncio.run(update_genshin())
-#     return 'Update smth ig!!'
-
-
-
-def upload_image(data):
-    print(data)
-    url = "https://fighter-programmer-uploader.hf.space/upload"
-    files = {'file': ('file', data, "image/png")}  
-    response = requests.post(url, files=files)
-    
-    if response.status_code != 200:
-        raise Exception(f"HTTP Error: {response.status_code}")
-    
-    try:
-        body = response.json()
-        if body["url"]:
-            return body["url"]
-        else:
-            raise Exception(f"Telegraph error: {body.get('error', 'Unknown error')}")
-    except (ValueError, KeyError, IndexError) as e:
-        raise Exception(f"Failed to parse response: {str(e)}")
-def process_image(dt):
-    with BytesIO() as byte_io:
-        dt.card.save(byte_io, "PNG")
-        byte_io.seek(0)
-        image_url = upload_image(byte_io)
+        # Get the system's temporary directory
+        temp_dir = tempfile.gettempdir()
+        temp_filename = os.path.join(temp_dir, f"{uuid.uuid4()}.jpg")
+        
+        # Download the image
+        async with httpx.AsyncClient() as client:
+            response = await client.get(url, timeout=timeout)
+            response.raise_for_status()
+            
+            # Save the image
+            with open(temp_filename, 'wb') as f:
+                f.write(response.content)
+        
+        # Open and process the image
+        img = Image.open(temp_filename)
+        
+        # Perform inference
+        with torch.no_grad():
+            inputs = processor(images=img, return_tensors="pt")
+            outputs = model(**inputs)
+            logits = outputs.logits
+        
+        # Calculate softmax probabilities
+        confidences = torch.softmax(logits, dim=-1)
+        
+        # Find the label with the highest confidence
+        max_confidence_id = confidences[0].argmax().item()
+        max_label = model.config.id2label[max_confidence_id]
+        max_confidence = confidences[0][max_confidence_id].item()
+        
+        # Clean up the temporary file
+        os.unlink(temp_filename)
         
         return {
-            "name": dt.name,
-            "url": image_url
+            "label": max_label,
+            "confidence": max_confidence
         }
-
-def process_images(result):
-    characters = []
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        # Execute image uploads in parallel
-        futures = [executor.submit(process_image, dt) for dt in result.card]
-        
-        for future in concurrent.futures.as_completed(futures):
-            try:
-                characters.append(future.result())
-            except Exception as e:
-                print(f"Error processing image: {e}")
     
-    return characters
-if __name__ == "__main__":
-    # data_dir = "/tmp/data"
-
-    # if not os.path.exists(data_dir):
-    #     os.makedirs(data_dir)
-
-    # data_dir = "/tmp/langs"
-
-    # if not os.path.exists(data_dir):
-    #     os.makedirs(data_dir)
-    # asyncio.run(update_genshin())
-    uvicorn.run("main:app", host="0.0.0.0", port=7860, workers=8, timeout_keep_alive=60000)
-# import enkacard
-# import enkacard.encbanner
-# import enkanetwork
-# import asyncio
-# import requests
-# from io import BytesIO
-# from enkanetwork import EnkaNetworkAPI
-# import os
-# from contextlib import asynccontextmanager
-
-# from fastapi import FastAPI
-
-# from fastapi.responses import (
-#     JSONResponse,
-# )  
-# from enkacard import enkatools
-
-# import uvicorn
-
-# enka_update = EnkaNetworkAPI()
-
-
-# async def update_genshin():
-#     try:
-#         async with enka_update:
-#             await enka_update.update_assets(path="/tmp", lang=["EN"])
-    
-
-#             # lang=["EN"], 
-#             print("Finished update")
-#     except Exception as e:
-#         print(f"Update failed: {e}")
-
-# async def card(id, designtype):
-#     async with enkacard.encbanner.ENC(uid = str(id)) as encard:
-#         return await encard.creat(template = (2 if str(designtype) == "2" else 1))
-    
-# @asynccontextmanager
-# async def lifespan(app: FastAPI):
-#     data_dir = "/tmp/data"
-
-#     if not os.path.exists(data_dir):
-#         os.makedirs(data_dir)
-
-#     data_dir = "/tmp/langs"
-
-#     if not os.path.exists(data_dir):
-#         os.makedirs(data_dir)
-#     await update_genshin()
-#     yield
-#     print("Goodbye")
-
-# app = FastAPI()
-# # app = FastAPI(lifespan=lifespan)
-
-# @app.get("/{id}")  # Correct route definition without prefix
-# async def characters(id: int, design: str = "1"):  # Use async and await
-#     try:
-#         characters = []
-#         result = await card(id, design)  # Use await instead of asyncio.run()
-
-#         for dt in result.card:
-#             with BytesIO() as byte_io:
-#                 dt.card.save(byte_io, "PNG")
-#                 byte_io.seek(0)
-#                 image_url = upload_image(byte_io)
-
-#                 characters.append({
-#                     "name": dt.name,
-#                     "url": image_url
-#                 })
-
-#         # Return valid JSON response using FastAPI's JSONResponse
-#         return JSONResponse(content={'response': characters})
-
-#     except enkanetwork.exception.VaildateUIDError:
-#         return JSONResponse(content={'error': 'Invalid UID. Please check your UID.'}, status_code=400)
-
-#     except enkacard.enc_error.ENCardError:
-#         return JSONResponse(content={'error': 'Enable display of the showcase in the game or add characters there.'}, status_code=400)
-
-#     except Exception as e:
-#         return JSONResponse(content={'error': 'UNKNOWN ERR: ' + str(e)}, status_code=500)
-
-# @app.get("/")
-# def hello_world():
-#     return 'AMERICA YA HALLO!!'
-
-# # @app.route("/update_char")
-# # def upload():
-#     # data_dir = "/tmp/data"
-
-#     # if not os.path.exists(data_dir):
-#     #     os.makedirs(data_dir)
-
-#     # data_dir = "/tmp/langs"
-
-#     # if not os.path.exists(data_dir):
-#     #     os.makedirs(data_dir)
-#     # asyncio.run(update_genshin())
-# #     return 'Update smth ig!!'
-
-
-
-# def upload_image(data):
-#     url = "https://telegra.ph/upload"
-#     files = {'file': ('file', data, "image/png")}  
-#     response = requests.post(url, files=files)
-    
-#     if response.status_code != 200:
-#         raise Exception(f"HTTP Error: {response.status_code}")
-    
-#     try:
-#         body = response.json()
-#         if isinstance(body, list) and 'src' in body[0]:
-#             return "https://telegra.ph" + body[0]['src']
-#         else:
-#             raise Exception(f"Telegraph error: {body.get('error', 'Unknown error')}")
-#     except (ValueError, KeyError, IndexError) as e:
-#         raise Exception(f"Failed to parse response: {str(e)}")
-    
-# if __name__ == "__main__":
-#     data_dir = "/tmp/data"
-
-#     if not os.path.exists(data_dir):
-#         os.makedirs(data_dir)
+    except httpx.RequestError as e:
+        raise HTTPException(status_code=400, detail=f"Error downloading image: {str(e)}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
 
-#     data_dir = "/tmp/langs"
+# Startup event to warm up the model
+@app.on_event("startup")
+async def startup_event():
+    # Perform a dummy inference to warm up the model
+    dummy_img = Image.new('RGB', (224, 224), color='red')
+    with torch.no_grad():
+        inputs = processor(images=dummy_img, return_tensors="pt")
+        model(**inputs)
 
-#     if not os.path.exists(data_dir):
-#         os.makedirs(data_dir)
-#     asyncio.run(update_genshin())
-#     uvicorn.run("main:app", host="0.0.0.0", port=7860, workers=8, timeout_keep_alive=600)
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)