修改跟踪算法cpu gpu切换错误

This commit is contained in:
2025-01-15 15:38:50 +08:00
parent 31db54fa37
commit ad850221c5
4 changed files with 16 additions and 11 deletions

View File

@ -122,6 +122,16 @@ def get_product_description(std_img, track_imgs):
return contrast_pair
def item_analysis(stream_dict):
track_imgs = stream_pipeline(stream_dict)
std_img = None
if stream_dict['goodsPic'] is not None:
response = requests.get(stream_dict['goodsPic'])
std_img = Image.open(BytesIO(response.content))
description_dict = get_product_description(std_img, track_imgs)
return description_dict
def main():
# sample input dict
stream_dict = {
@ -137,13 +147,8 @@ def main():
"goodsSpec" : "405g"
}
track_imgs = stream_pipeline(stream_dict)
std_img = None
if stream_dict['goodsPic'] is not None:
response = requests.get(stream_dict['goodsPic'])
std_img = Image.open(BytesIO(response.content))
description_dict = get_product_description(std_img, track_imgs)
print(description_dict)
result = item_analysis(stream_dict)
print(result)
if __name__ == "__main__":
main()

View File

@ -108,7 +108,7 @@ class FeatsInterface:
# patch = patch.to(device=self.device).half()
if str(self.device) != "cpu":
patch = patch.to(device=self.device).half()
patch = patch.to(device=self.device)
else:
patch = patch.to(device=self.device)

View File

@ -217,7 +217,7 @@ def yolo_resnet_tracker(
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
# det = det.cpu().numpy()
det = det.cpu().numpy()
## ================================================================ writed by WQG
'''tracks: [x1, y1, x2, y2, track_id, score, cls, frame_index, box_index]
0 1 2 3 4 5 6 7 8

View File

@ -153,8 +153,8 @@ class doBackTracks(doTracks):
hand_ious = []
hboxes = np.empty(shape=(0, 9), dtype = np.float)
gboxes = np.empty(shape=(0, 9), dtype = np.float)
hboxes = np.empty(shape=(0, 9), dtype = float)
gboxes = np.empty(shape=(0, 9), dtype = float)
# start, end 为索引值,需要 start:(end+1)