routers

Author: JuliaBars

app.py renamed to bll.py

@@ -11,44 +11,48 @@
 
 from models.sample_model.labels import cars_models
 
-# Initialize the models
-model_sample_model = YOLO("./models/sample_model/yolov8n.pt")
+# Initialize the model
+model_sample_model = YOLO('./models/sample_model/yolov8n.pt')
 
 
+# region Common func
+
 def get_image_from_bytes(binary_image: bytes) -> Image:
-    """ Convert image from bytes to PIL RGB format. """
-    input_image = Image.open(io.BytesIO(binary_image)).convert("RGB")
+    ''' Convert image from bytes to PIL RGB format. '''
+    input_image = Image.open(io.BytesIO(binary_image)).convert('RGB')
     return input_image
 
 
 def get_bytes_from_image(image: Image) -> bytes:
-    """ Convert PIL image to Bytes. """
+    ''' Convert PIL image to Bytes. '''
     return_image = io.BytesIO()
     image.save(return_image, format='JPEG', quality=100)
     return_image.seek(0)
     return return_image
 
+# endregion
+
+# region Yolo func
+
 
 def transform_predict_to_df(results: list, labeles_dict: dict) -> pd.DataFrame:
-    """
-    Transform predict from yolov8 (torch.Tensor) to pandas DataFrame.   
-    """
+    ''' Transform predict from yolov8 (torch.Tensor) to pandas DataFrame. '''
     predict_bbox = pd.DataFrame(
-        results[0].to("cpu").numpy().boxes.xyxy, columns=[
-            'xmin', 'ymin', 'xmax', 'ymax']
+        results[0].to('cpu').numpy().boxes.xyxy, columns=[
+            'xmin', "ymin", 'xmax', 'ymax']
         )
-    predict_bbox['confidence'] = results[0].to("cpu").numpy().boxes.conf
+    predict_bbox['confidence'] = results[0].to('cpu').numpy().boxes.conf
     predict_bbox['class'] = (
-        results[0].to("cpu").numpy().boxes.cls).astype(int)
-    predict_bbox['name'] = predict_bbox["class"].replace(labeles_dict)
+        results[0].to('cpu').numpy().boxes.cls).astype(int)
+    predict_bbox['name'] = predict_bbox['class'].replace(labeles_dict)
     return predict_bbox
 
 
 def get_model_predict(
         model: YOLO, input_image: Image,
         save: bool = False, image_size: int = 1248,
         conf: float = 0.5, augment: bool = False) -> pd.DataFrame:
-    """ Get the predictions of a model on an input image. """
+    ''' Get the predictions of a model on an input image. '''
     predictions = model.predict(
                         imgsz=image_size,
                         source=input_image,
@@ -63,10 +67,8 @@ def get_model_predict(
     return predictions
 
 
-# ----------------------- BBOX Func -----------------------
-
 def add_bboxs_on_img(image: Image, predict: pd.DataFrame()) -> Image:
-    """ Add a bounding box on the image. """
+    ''' Add a bounding box on the image. '''
     annotator = Annotator(np.array(image))
     predict = predict.sort_values(by=['xmin'], ascending=True)
     for i, row in predict.iterrows():
@@ -77,7 +79,7 @@ def add_bboxs_on_img(image: Image, predict: pd.DataFrame()) -> Image:
 
 
 def detect_sample_model(input_image: Image) -> pd.DataFrame:
-    """ Predict from sample_model. Base on YoloV8. """
+    ''' Predict from sample_model. Base on YoloV8. '''
     predict = get_model_predict(
         model=model_sample_model,
         input_image=input_image,
@@ -88,14 +90,16 @@ def detect_sample_model(input_image: Image) -> pd.DataFrame:
     )
     return predict
 
+# endregion
+
+# region Car classification func
 
-#------------------------MODELS CLASSIFICATION------------------------
 
 class YOLOSegmentation:
-    """
+    '''
     Useful class to get bboxes, classes, segmentations, scores in correct
     format to pass them to cv2 image processed functions.
-    """
+    '''
     def __init__(self, model_path):
         self.model = YOLO(model_path)
 
@@ -117,7 +121,7 @@ def detect(self, img):
         return bboxes, class_ids, segmentation_contours_idx, scores
 
 
-#model for segmentation
+# Initialize model for segmentation
 model_for_classify = YOLOSegmentation("./models/sample_model/yolov8m-seg.pt")
 
 # for model classification project we will segment only cars, busses, truckes
@@ -129,7 +133,7 @@ def detect(self, img):
 
 
 def extract_segment_image(img, segmentator=model_for_classify) -> Image:
-    """ Extract segmented and cropped car, truck, bus and return as PIL"""
+    ''' Extract segmented and cropped car, truck, bus and return as PIL. '''
     open_cv_image = np.array(img)
     open_cv_image = open_cv_image[:, :, ::-1].copy()
 
@@ -151,16 +155,16 @@ def extract_segment_image(img, segmentator=model_for_classify) -> Image:
 
 
 def load_model():
-    """ Load and evaluate saved model. """
+    ''' Load and evaluate saved model. '''
     model = torch.load(
         './models/sample_model/model_mob_netv3_79_perc.pth',
         map_location=torch.device('cpu'))
     model.eval()
     return model
 
 
-def image_to_tensor(cv2_img: np.ndarray) ->  torch.Tensor:
-    """ Converting cv2 output to torch tensor. """
+def image_to_tensor(cv2_img: np.ndarray) -> torch.Tensor:
+    ''' Converting cv2 output to torch tensor. '''
     image = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB)
     transform = transforms.Compose([
         transforms.ToTensor(),
@@ -180,6 +184,7 @@ def draw_predictions_on_image(img: Image, label: str) -> Image:
 
 
 def predict_brand_and_model(img, segmentator=model_for_classify):
+    ''' Predict car's brand and model using PyTorch model. '''
     open_cv_image = np.array(img)
     open_cv_image = open_cv_image[:, :, ::-1].copy()
 
@@ -201,7 +206,7 @@ def predict_brand_and_model(img, segmentator=model_for_classify):
     answer = predict.argmax(-1)
     name = cars_models.get(answer.item()).split('_')
     name = f'Brand: {name[0]}, model: {name[1]}'
-    im_rgb = cv2.cvtColor(crop_img, cv2.COLOR_BGR2RGB)
-    # PIL_im = Image.fromarray(im_rgb)
-    PIL_im = draw_predictions_on_image(img, name)
-    return PIL_im
+    im_PIL = draw_predictions_on_image(img, name)
+    return im_PIL
+
+# endregion

logger_cfg.py

@@ -0,0 +1,16 @@
+config = {
+    "handlers": [
+        {"sink": "log.log", 
+         "format": "<green>{time:HH:mm:ss}</green> | <level>{message}</level>",
+         'rotation': "1 MB", 'level': "DEBUG", 'compression': "zip"},
+    ],
+}
+
+
+# logger.remove()
+# logger.add(
+#     colorize=True,
+#     format="<green>{time:HH:mm:ss}</green> | <level>{message}</level>",
+#     level=10,
+# )
+# logger.add("log.log", rotation="1 MB", level="DEBUG", compression="zip")

main.py

@@ -1,50 +1,22 @@
 import json
-import pandas as pd
-from PIL import Image
-from loguru import logger
-import sys
 
-from fastapi import FastAPI, File, status
+from fastapi import FastAPI
 from fastapi.responses import RedirectResponse
-from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.exceptions import HTTPException
 
-from app import get_image_from_bytes
-from app import detect_sample_model
-from app import add_bboxs_on_img
-from app import get_bytes_from_image
-from app import extract_segment_image
-from app import predict_brand_and_model
+from routers import check, classify, try_yolo
 
 
-from models.sample_model.labels import cars_models
-from models.sample_model.yolo_labels import yolo_classes
+# region FastAPI SetUp
 
-
-logger.remove()
-logger.add(
-    sys.stderr,
-    colorize=True,
-    format="<green>{time:HH:mm:ss}</green> | <level>{message}</level>",
-    level=10,
-)
-logger.add("log.log", rotation="1 MB", level="DEBUG", compression="zip")
-
-# ------------------ FastAPI Setup ------------------
-
-# title
 app = FastAPI(
-    title="Object Detection FastAPI Template",
+    title="Car models Classification FastAPI",
     description="""Obtain object value out of image
-                    and return image and json result""",
-    version="2023.1.31",
+                    and return image with label""",
+    version="2023.7.16",
 )
 
-# This function is needed if you want to allow client requests
-# from specific domains (specified in the origins argument)
-# to access resources from the FastAPI server,
-# and the client and server are hosted on different domains.
+
 origins = [
     "http://localhost",
     "http://localhost:8008",
@@ -62,144 +34,23 @@
 
 @app.on_event("startup")
 def save_openapi_json():
-    '''This function is used to save the OpenAPI documentation
+    ''' This function is used to save the OpenAPI documentation
     data of the FastAPI application to a JSON file.
     The purpose of saving the OpenAPI documentation data is to
     which can be used for documentation purposes or
-    to generate client libraries. It is not necessarily needed,
-    but can be helpful in certain scenarios.'''
+    to generate client libraries. '''
     openapi_data = app.openapi()
-    # Change "openapi.json" to desired filename
     with open("openapi.json", "w") as file:
         json.dump(openapi_data, file)
 
+# endregion
+
 
-# redirect
 @app.get("/", include_in_schema=False)
 async def redirect():
     return RedirectResponse("/docs")
 
 
-@app.get('/healthcheck', status_code=status.HTTP_200_OK)
-def perform_healthcheck():
-    '''
-    It basically sends a GET request to the route & hopes to get a "200"
-    response code. Failing to return a 200 response code just enables
-    the GitHub Actions to rollback to the last version the project was
-    found in a "working condition". It acts as a last line of defense in
-    case something goes south.
-    Additionally, it also returns a JSON response in the form of:
-    {
-        'healtcheck': 'Everything OK!'
-    }
-    '''
-    return {'healthcheck': 'Everything OK!'}
-
-
-# ------------------ Support Func ------------------
-
-def crop_image_by_predict(
-        image: Image,
-        predict: pd.DataFrame(),
-        crop_class_name: str,) -> Image:
-    """Crop an image based on the detection of a certain object in the image.
-
-    Args:
-        image: Image to be cropped.
-        predict (pd.DataFrame): Dataframe containing the prediction results
-        of object detection model.
-        crop_class_name (str, optional): The name of the object class to crop
-        the image by. if not provided, function returns the first object found
-        in the image.
-
-    Returns:
-        Image: Cropped image or None
-    """
-    crop_predicts = predict[(predict['name'] == crop_class_name)]
-
-    if crop_predicts.empty:
-        raise HTTPException(
-            status_code=400, detail=f"{crop_class_name} not found in photo")
-
-    # if there are several detections, choose the one with more confidence
-    if len(crop_predicts) > 1:
-        crop_predicts = crop_predicts.sort_values(
-            by=['confidence'], ascending=False)
-
-    crop_bbox = crop_predicts[['xmin', 'ymin', 'xmax', 'ymax']].iloc[0].values
-    img_crop = image.crop(crop_bbox)
-    return img_crop
-
-
-# ------------------ MAIN Func ------------------
-
-@app.get('/available_models', status_code=status.HTTP_200_OK)
-def get_all_models():
-    '''
-    To check the list of cars' models, which are avalable for recognition.
-    '''
-    return cars_models
-
-@app.get('/yolo_classes', status_code=status.HTTP_200_OK)
-def get_yolo_classes():
-    '''
-    To check the list of cars' models, which are avalable for recognition.
-    '''
-    return yolo_classes
-
-@app.post("/object_detection_with_yolo")
-def object_detection_with_yolo(file: bytes = File(...)):
-    """
-    Object Detection from an image plot bbox on image. Using Yolo8.
-
-    Args:
-        file (bytes): The image file in bytes format.
-    Returns:
-        Image: Image in bytes with bbox annotations.
-    """
-    # get image from bytes
-    input_image = get_image_from_bytes(file)
-
-    # model predict
-    predict = detect_sample_model(input_image)
-
-    # add bbox on image
-    final_image = add_bboxs_on_img(image=input_image, predict=predict)
-
-    # return image in bytes format
-    return StreamingResponse(
-        content=get_bytes_from_image(final_image), media_type="image/jpeg")
-
-
-@logger.catch
-@app.post("/car_brand_model_classification")
-def car_brand_model_classification(file: bytes = File(...)):
-    """
-    Object Detection from an image plot bbox on image. Using Yolo8.
-
-    Args:
-        file (bytes): The image file in bytes format.
-    Returns:
-        Image: Image in bytes with bbox annotations.
-    """
-    input_image = get_image_from_bytes(file)
-    predict = predict_brand_and_model(input_image)
-    return StreamingResponse(
-        content=get_bytes_from_image(predict), media_type="image/jpeg")
-
-
-@app.post("/car_model_segment_and_crop")
-def car_model_segment_and_crop(file: bytes = File(...)):
-    """
-    Object Detection from an image plot bbox on image. Using Yolo8.
-
-    Args:
-        file (bytes): The image file in bytes format.
-    Returns:
-        Image: Image in bytes with bbox annotations.
-    """
-    input_image = get_image_from_bytes(file)
-    predict = extract_segment_image(input_image)
-
-    return StreamingResponse(
-        content=get_bytes_from_image(predict), media_type="image/jpeg")
+app.include_router(check.router)
+app.include_router(classify.router)
+app.include_router(try_yolo.router)