how to Load YOLOv7 using torch.hub - pytorch

Question

how to load yolov7 model using torch.hub for make prediction
I directly use torch.hub.load method of yolov5 but it didn't work

Answer 1

The answer is :
First method
!# Download YOLOv7 code
!git clone https://github.com/WongKinYiu/yolov7
%cd yolov7
from hubconf import custom
model = custom(path_or_model='yolov7.pt') # custom example
# model = create(name='yolov7', pretrained=True, channels=3, classes=80, autoshape=True) # pretrained example
# Verify inference
import numpy as np
from PIL import Image
imgs = [np.zeros((640, 480, 3))]
results = model(imgs) # batched inference
results.print()
results.save()
Second method
Load YOLOv7 using torch.hub
!# Download YOLOv7 code
!git clone https://github.com/WongKinYiu/yolov7
%cd yolov7
from pathlib import Path
import torch
from models.yolo import Model
from utils.general import check_requirements, set_logging
from utils.google_utils import attempt_download
from utils.torch_utils import select_device
dependencies = ['torch', 'yaml']
check_requirements(Path("/content/yolov7/").parent / 'requirements.txt', exclude=('pycocotools', 'thop'))
set_logging()
def custom(path_or_model='path/to/model.pt', autoshape=True):
"""custom mode
Arguments (3 options):
path_or_model (str): 'path/to/model.pt'
path_or_model (dict): torch.load('path/to/model.pt')
path_or_model (nn.Module): torch.load('path/to/model.pt')['model']
Returns:
pytorch model
"""
model = torch.load(path_or_model, map_location=torch.device('cpu')) if isinstance(path_or_model, str) else path_or_model # load checkpoint
if isinstance(model, dict):
model = model['ema' if model.get('ema') else 'model'] # load model
hub_model = Model(model.yaml).to(next(model.parameters()).device) # create
hub_model.load_state_dict(model.float().state_dict()) # load state_dict
hub_model.names = model.names # class names
if autoshape:
hub_model = hub_model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS
device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available
return hub_model.to(device)
model = custom(path_or_model='yolov7.pt') # custom example
# model = create(name='yolov7', pretrained=True, channels=3, classes=80, autoshape=True) # pretrained example
# Verify inference
import numpy as np
from PIL import Image
imgs = [np.zeros((640, 480, 3))]
results = model(imgs) # batched inference
results.print()
results.save()
df_prediction = results.pandas().xyxy
df_prediction
full code on google colab

Answer 2

The Yolov7 model was fine-tuned on custom dataset with this repo: https://github.com/WongKinYiu/yolov7
import torch
import os
import cv2
# Load fine-tuned custom model
model = torch.hub.load('WongKinYiu/yolov7', 'custom', '/path/to/custom_model.pt',
force_reload=True, trust_repo=True)
# Declaring some variables
TABLE_CONFIDENCE = 0.50
CELL_CONFIDENCE = 0.50
OUTPUT_DIR = 'output'
# Bounding Boxes color scheme
ALPHA = 0.2
TABLE_BORDER = (0, 0, 255)
CELL_FILL = (0, 0, 200)
CELL_BORDER = (0, 0, 255)
os.makedirs(OUTPUT_DIR, exist_ok=True)
# Run the Inference and draw predicted bboxes
results = model(image_path)
df = results.pandas().xyxy[0]
table_bboxes = []
cell_bboxes = []
for _, row in df.iterrows():
if row['class'] == 0 and row['confidence'] > TABLE_CONFIDENCE:
table_bboxes.append([int(row['xmin']), int(row['ymin']),
int(row['xmax']), int(row['ymax'])])
if row['class'] == 1 and row['confidence'] > CELL_CONFIDENCE:
cell_bboxes.append([int(row['xmin']), int(row['ymin']),
int(row['xmax']), int(row['ymax'])])
image = cv2.imread(image_path)
overlay = image.copy()
for table_bbox in table_bboxes:
cv2.rectangle(image, (table_bbox[0], table_bbox[1]),
(table_bbox[2], table_bbox[3]), TABLE_BORDER, 1)
for cell_bbox in cell_bboxes:
cv2.rectangle(overlay, (cell_bbox[0], cell_bbox[1]),
(cell_bbox[2], cell_bbox[3]), CELL_FILL, -1)
cv2.rectangle(image, (cell_bbox[0], cell_bbox[1]),
(cell_bbox[2], cell_bbox[3]), CELL_BORDER, 1)
image_new = cv2.addWeighted(overlay, ALPHA, image, 1-ALPHA, 0)
image_filename = image_path.split('/')[-1]
cv2.imwrite(f'{OUTPUT_DIR}/{image_filename}', image_new)