File size: 5,550 Bytes
7336c79 dc39c99 7336c79 dc39c99 7336c79 e5c469f 7336c79 e5c469f 7336c79 e5c469f 7336c79 e2a86b4 7336c79 dc39c99 7336c79 224d736 7336c79 dc39c99 4f45aec dc39c99 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 |
from fastapi import FastAPI, File, UploadFile, Response, HTTPException
from fastapi.responses import JSONResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
import io
import sqlite3
from pydantic import BaseModel, EmailStr
from pathlib import Path
from model import YOLOModel
import shutil
yolo = YOLOModel()
UPLOAD_FOLDER = Path("./uploads")
UPLOAD_FOLDER.mkdir(exist_ok=True)
app = FastAPI()
cropped_images_dir = "cropped_images"
# Initialize SQLite database
def init_db():
conn = sqlite3.connect('users.db')
c = conn.cursor()
c.execute('''
CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
firstName TEXT NOT NULL,
lastName TEXT NOT NULL,
country TEXT,
number TEXT, -- Phone number stored as TEXT to allow various formats
email TEXT UNIQUE NOT NULL, -- Email should be unique and non-null
password TEXT NOT NULL -- Password will be stored as a string (hashed ideally)
)
''')
conn.commit()
conn.close()
init_db()
class UserSignup(BaseModel):
firstName: str
lastName: str
country: str
number: str
email: EmailStr
password: str
class UserLogin(BaseModel):
email: str
password: str
@app.post("/signup")
async def signup(user_data: UserSignup):
try:
conn = sqlite3.connect('users.db')
c = conn.cursor()
# Check if user already exists
c.execute("SELECT * FROM users WHERE email = ?", (user_data.email,))
if c.fetchone():
raise HTTPException(status_code=400, detail="Email already registered")
# Insert new user
c.execute("""
INSERT INTO users (firstName, lastName, country, number, email, password)
VALUES (?, ?, ?, ?, ?, ?)
""", (user_data.firstName, user_data.lastName, user_data.country, user_data.number, user_data.email, user_data.password))
conn.commit()
conn.close()
return {"message": "User registered successfully", "email": user_data.email}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/login")
async def login(user_data: UserLogin):
try:
conn = sqlite3.connect('users.db')
c = conn.cursor()
# Find user
c.execute("SELECT * FROM users WHERE email = ? AND password = ?",
(user_data.email, user_data.password))
user = c.fetchone()
conn.close()
if not user:
raise HTTPException(status_code=401, detail="Invalid credentials")
return {
"message": "Login successful",
"user": {
"firstName": user[1],
"lastName": user[2],
"email": user[3]
}
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/users")
async def get_users():
try:
conn = sqlite3.connect('users.db')
conn.row_factory = sqlite3.Row # This makes rows behave like dictionaries
c = conn.cursor()
c.execute("SELECT * FROM users")
rows = c.fetchall()
conn.close()
# Convert rows to a list of dictionaries
users = [dict(row) for row in rows]
return users
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/upload")
async def upload_image(image: UploadFile = File(...)):
# print(f'\n\t\tUPLOADED!!!!')
try:
file_path = UPLOAD_FOLDER / image.filename
with file_path.open("wb") as buffer:
shutil.copyfileobj(image.file, buffer)
# print(f'Starting to pass into model, {file_path}')
# Perform YOLO inference
predictions = yolo.predict(str(file_path))
print(f'\n\n\n{predictions}\n\n\ \n\t\t\t\tare predictions')
# Clean up uploaded file
file_path.unlink() # Remove file after processing
return JSONResponse(content={"items": predictions})
except Exception as e:
return JSONResponse(content={"error": str(e)}, status_code=500)
def cleanup_images(directory: str):
"""Remove all images in the directory."""
for file in Path(directory).glob("*"):
file.unlink()
# @app.post("/upload")
# async def upload_image(image: UploadFile = File(...)):
# # print(f'\n\t\tUPLOADED!!!!')
# try:
# file_path = UPLOAD_FOLDER / image.filename
# with file_path.open("wb") as buffer:
# shutil.copyfileobj(image.file, buffer)
# # print(f'Starting to pass into model, {file_path}')
# # Perform YOLO inference
# predictions = yolo.predict(str(file_path))
# print(f'\n\n\n{predictions}\n\n\ \n\t\t\t\tare predictions')
# # Clean up uploaded file
# file_path.unlink() # Remove file after processing
# return JSONResponse(content={"items": predictions})
# except Exception as e:
# return JSONResponse(content={"error": str(e)}, status_code=500)
# code to accept the localhost to get images from
app.add_middleware(
CORSMiddleware,
allow_origins=["http://192.168.56.1:3000", "http://192.168.56.1:3001", "https://recognizethis.netlify.app/", "*"],
allow_methods=["*"],
allow_headers=["*"],
)
if __name__ == "__main__":
import uvicorn
uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)
|