Ercik commited on
Commit
8de87ce
verified
1 Parent(s): 5f70109

Create app/main.py

Browse files
Files changed (1) hide show
  1. 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")