Spaces:
Sleeping
Sleeping
Create app/main.py
Browse files- app/main.py +122 -0
app/main.py
ADDED
|
@@ -0,0 +1,122 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from fastapi import FastAPI, Depends, HTTPException, status, File, UploadFile, Response
|
| 2 |
+
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
|
| 3 |
+
from jose import JWTError, jwt
|
| 4 |
+
from passlib.context import CryptContext
|
| 5 |
+
from pydantic import BaseModel
|
| 6 |
+
from datetime import datetime, timedelta
|
| 7 |
+
from io import BytesIO
|
| 8 |
+
from PIL import Image, ImageDraw
|
| 9 |
+
import numpy as np
|
| 10 |
+
from ultralytics import YOLO
|
| 11 |
+
|
| 12 |
+
# Configuraci贸n JWT
|
| 13 |
+
SECRET_KEY = "clave-ultra-secreta"
|
| 14 |
+
ALGORITHM = "HS256"
|
| 15 |
+
ACCESS_TOKEN_EXPIRE_MINUTES = 30
|
| 16 |
+
|
| 17 |
+
# Hash de contrase帽as
|
| 18 |
+
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
|
| 19 |
+
|
| 20 |
+
# Base de datos falsa
|
| 21 |
+
fake_users_db = {
|
| 22 |
+
"uriel": {
|
| 23 |
+
"username": "uriel",
|
| 24 |
+
"full_name": "Uriel Escalona",
|
| 25 |
+
"email": "uriel@example.com",
|
| 26 |
+
"hashed_password": pwd_context.hash("1234"),
|
| 27 |
+
"disabled": False
|
| 28 |
+
}
|
| 29 |
+
}
|
| 30 |
+
|
| 31 |
+
# Modelos
|
| 32 |
+
class Token(BaseModel):
|
| 33 |
+
access_token: str
|
| 34 |
+
token_type: str
|
| 35 |
+
|
| 36 |
+
class TokenData(BaseModel):
|
| 37 |
+
username: str | None = None
|
| 38 |
+
|
| 39 |
+
class User(BaseModel):
|
| 40 |
+
username: str
|
| 41 |
+
full_name: str
|
| 42 |
+
email: str
|
| 43 |
+
disabled: bool | None = None
|
| 44 |
+
|
| 45 |
+
class UserInDB(User):
|
| 46 |
+
hashed_password: str
|
| 47 |
+
|
| 48 |
+
# Autenticaci贸n
|
| 49 |
+
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
|
| 50 |
+
|
| 51 |
+
app = FastAPI()
|
| 52 |
+
model = YOLO("yolov8n.pt")
|
| 53 |
+
|
| 54 |
+
def verify_password(plain_password, hashed_password):
|
| 55 |
+
return pwd_context.verify(plain_password, hashed_password)
|
| 56 |
+
|
| 57 |
+
def get_user(db, username: str):
|
| 58 |
+
if username in db:
|
| 59 |
+
user_dict = db[username]
|
| 60 |
+
return UserInDB(**user_dict)
|
| 61 |
+
|
| 62 |
+
def authenticate_user(db, username: str, password: str):
|
| 63 |
+
user = get_user(db, username)
|
| 64 |
+
if not user or not verify_password(password, user.hashed_password):
|
| 65 |
+
return False
|
| 66 |
+
return user
|
| 67 |
+
|
| 68 |
+
def create_access_token(data: dict, expires_delta: timedelta | None = None):
|
| 69 |
+
to_encode = data.copy()
|
| 70 |
+
expire = datetime.utcnow() + (expires_delta or timedelta(minutes=15))
|
| 71 |
+
to_encode.update({"exp": expire})
|
| 72 |
+
return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
|
| 73 |
+
|
| 74 |
+
async def get_current_user(token: str = Depends(oauth2_scheme)):
|
| 75 |
+
credentials_exception = HTTPException(
|
| 76 |
+
status_code=401,
|
| 77 |
+
detail="No autorizado",
|
| 78 |
+
headers={"WWW-Authenticate": "Bearer"},
|
| 79 |
+
)
|
| 80 |
+
try:
|
| 81 |
+
payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
|
| 82 |
+
username: str = payload.get("sub")
|
| 83 |
+
if username is None:
|
| 84 |
+
raise credentials_exception
|
| 85 |
+
user = get_user(fake_users_db, username)
|
| 86 |
+
if user is None:
|
| 87 |
+
raise credentials_exception
|
| 88 |
+
return user
|
| 89 |
+
except JWTError:
|
| 90 |
+
raise credentials_exception
|
| 91 |
+
|
| 92 |
+
# Endpoint para login
|
| 93 |
+
@app.post("/token", response_model=Token)
|
| 94 |
+
async def login(form_data: OAuth2PasswordRequestForm = Depends()):
|
| 95 |
+
user = authenticate_user(fake_users_db, form_data.username, form_data.password)
|
| 96 |
+
if not user:
|
| 97 |
+
raise HTTPException(status_code=401, detail="Credenciales inv谩lidas")
|
| 98 |
+
access_token = create_access_token(data={"sub": user.username}, expires_delta=timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES))
|
| 99 |
+
return {"access_token": access_token, "token_type": "bearer"}
|
| 100 |
+
|
| 101 |
+
# Endpoint protegido: recibe imagen y devuelve con cajas
|
| 102 |
+
@app.post("/detectar/")
|
| 103 |
+
async def detectar_yolo(file: UploadFile = File(...), current_user: User = Depends(get_current_user)):
|
| 104 |
+
contents = await file.read()
|
| 105 |
+
image = Image.open(BytesIO(contents)).convert("RGB")
|
| 106 |
+
img_array = np.array(image)
|
| 107 |
+
|
| 108 |
+
results = model.predict(img_array)
|
| 109 |
+
draw = ImageDraw.Draw(image)
|
| 110 |
+
|
| 111 |
+
for r in results:
|
| 112 |
+
for box in r.boxes:
|
| 113 |
+
cls = int(box.cls[0])
|
| 114 |
+
label = model.names[cls]
|
| 115 |
+
bbox = box.xyxy[0].tolist()
|
| 116 |
+
draw.rectangle(bbox, outline="red", width=2)
|
| 117 |
+
draw.text((bbox[0], bbox[1]), label, fill="white")
|
| 118 |
+
|
| 119 |
+
output = BytesIO()
|
| 120 |
+
image.save(output, format="PNG")
|
| 121 |
+
output.seek(0)
|
| 122 |
+
return Response(content=output.read(), media_type="image/png")
|