Skip to content

Commit

Permalink
添加了对象描述gpt生成的功能，支持openai-api,本地部署chatglm3-6b测试可用
Browse files Browse the repository at this point in the history
  • Loading branch information
@sungatetop
sungatetop committed Jan 17, 2024
1 parent 1edd6cc commit cf3ca4d
Show file tree
Hide file tree
Showing 5 changed files with 142 additions and 44 deletions.
1 change: 0 additions & 1 deletion anylabeling/app.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -206,7 +206,6 @@ def main():
output_file=output_file,
output_dir=output_dir,
)

if reset_config:
logger.info("Resetting Qt config: %s", win.settings.fileName())
win.settings.clear()
Expand Down
6 changes: 5 additions & 1 deletion anylabeling/configs/anylabeling_config.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,11 @@ file_dock:
closable: false
movable: false
floatable: false

gpt_dock:
show: true
closable: false
movable: false
floatable: true
# label_dialog
show_label_text_field: true
label_completion: startswith
Expand Down
93 changes: 90 additions & 3 deletions anylabeling/views/labeling/label_widget.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import webbrowser
from difflib import SequenceMatcher

import openai
import darkdetect
import imgviz
import natsort
Expand All @@ -28,6 +29,8 @@
QMessageBox,
QProgressDialog,
QScrollArea,
QLineEdit,
QPushButton
)

from anylabeling.services.auto_labeling.types import AutoLabelingMode
Expand Down Expand Up @@ -1262,10 +1265,55 @@ def __init__(
"}"
)
self.shape_text_edit = QPlainTextEdit()
#添加gpt自动生成描述
self.gpt_dock = QtWidgets.QDockWidget(self.tr("GPT Text Generator"), self)
self.gpt_dock.setObjectName("GPT")
self.gpt_api_edit = QLineEdit()
self.gpt_api_edit.setPlaceholderText("Enter GPT API Endpoint")
self.gpt_token_edit = QLineEdit()
self.gpt_token_edit.setPlaceholderText("Enter GPT Token")
self.gpt_model_edit = QLineEdit()
self.gpt_model_edit.setPlaceholderText("Enter GPT Model Name")
self.gpt_role_edit = QLineEdit()
self.gpt_role_edit.setPlaceholderText("Enter GPT Role Description")
self.shape_prompt_edit = QLineEdit()
self.shape_prompt_edit.setPlaceholderText("Enter Shape Prompt")
# Creating a button for triggering text generation
self.text_generate_pb = QPushButton("Generate Text")

self.prev_image_pb = QPushButton("PreImage")
self.next_image_pb = QPushButton("NextImage")
h_layout = QHBoxLayout()
h_layout.addWidget(self.prev_image_pb)
h_layout.addWidget(self.text_generate_pb)
h_layout.addWidget(self.next_image_pb)

gpt_layout = QtWidgets.QVBoxLayout()
gpt_layout.setContentsMargins(0, 0, 0, 0)
gpt_layout.setSpacing(0)
gpt_layout.addWidget(self.gpt_api_edit)
gpt_layout.addWidget(self.gpt_token_edit)
gpt_layout.addWidget(self.gpt_model_edit)
gpt_layout.addWidget(self.gpt_role_edit)
gpt_layout.addWidget(self.shape_prompt_edit)
gpt_layout.addWidget(self.shape_text_edit)
gpt_layout.addLayout(h_layout)

#gpt_layout.addWidget(self.text_generate_pb)

gpt_widget=QtWidgets.QWidget()
gpt_widget.setLayout(gpt_layout)
self.gpt_dock.setWidget(gpt_widget)
self.gpt_dock.setStyleSheet(dock_title_style)
self.text_generate_pb.clicked.connect(self.shape_text_generate)
self.prev_image_pb.clicked.connect(self.open_prev_image)
self.next_image_pb.clicked.connect(self.open_next_image)

right_sidebar_layout.addWidget(
self.shape_text_label, 0, Qt.AlignCenter
)
right_sidebar_layout.addWidget(self.shape_text_edit)
#right_sidebar_layout.addWidget(self.shape_text_edit)
right_sidebar_layout.addWidget(self.gpt_dock)
right_sidebar_layout.addWidget(self.flag_dock)
right_sidebar_layout.addWidget(self.label_dock)
right_sidebar_layout.addWidget(self.label_filter_combobox)
Expand All @@ -1290,7 +1338,7 @@ def __init__(
self.shape_dock.setFeatures(
self.shape_dock.features() & rev_dock_features
)

self.shape_text_edit.textChanged.connect(self.shape_text_changed)

layout.addItem(right_sidebar_layout)
Expand Down Expand Up @@ -2023,7 +2071,46 @@ def add_label(self, shape):
label_list_item.setText("{}".format(html.escape(text)))
label_list_item.setBackground(QtGui.QColor(*color, LABEL_OPACITY))
self.update_combo_box()

#GPT text Generator
def shape_text_generate(self):
api_key=self.gpt_token_edit.text().strip()
api_endpoint = self.gpt_api_edit.text().strip()
model_name=self.gpt_model_edit.text().strip()
gpt_role=self.gpt_role_edit.text().strip()
prompt = self.shape_prompt_edit.text().strip()
#需要加一个判断是否使用gpt4，多模态的
print(api_endpoint,api_key,model_name,gpt_role,prompt,self.filename)
if api_key is None:
openai.api_key="EMPTY"
if api_endpoint is None:
QMessageBox.warning(self,"GPT warning","api is empty!")
return
if model_name is None:
QMessageBox.warning(self,"GPT warning","model_name is empty!")
return
if gpt_role is None:
gpt_role="you are a wonderful assistant for user"
if prompt is None:
QMessageBox.warning(self,"GPT warning","prompt is empty!")
return
messages=[
{"role":"system","content":gpt_role},
{"role":"user","content":prompt}
]
#如果使用gpt4V或其他多模态模型,openai版本需要更新,将图片一起打包即可，以下会有不同！
openai.api_key=api_key
openai.api_base=api_endpoint
try:
completion = openai.ChatCompletion.create(model=model_name,
messages=messages,
max_tokens=200)
text=str(completion.choices[0].message['content'])
self.set_text_editing(True)
self.shape_text_edit.setPlainText(text)
except Exception as e:
QMessageBox.warning(self,"GPT warning",f"Error: {str(e)}")


def shape_text_changed(self):
description = self.shape_text_edit.toPlainText()
if self.canvas.current is not None:
Expand Down
83 changes: 45 additions & 38 deletions anylabeling/views/labeling/widgets/canvas.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1111,45 +1111,52 @@ def paintEvent(self, event): # noqa: C901
drawing_shape.add_point(self.line[1])
drawing_shape.fill = True
drawing_shape.paint(p)

#文字描述显示有问题
# Draw texts
if self.show_texts:
p.setFont(
QtGui.QFont(
"Arial", int(max(6.0, int(round(8.0 / Shape.scale))))
)
)
pen = QtGui.QPen(QtGui.QColor("#00FF00"), 8, Qt.SolidLine)
p.setPen(pen)
for shape in self.shapes:
description = shape.description
if description:
bbox = shape.bounding_rect()
fm = QtGui.QFontMetrics(p.font())
rect = fm.boundingRect(description)
p.fillRect(
rect.x() + bbox.x(),
rect.y() + bbox.y(),
rect.width(),
rect.height(),
QtGui.QColor("#00FF00"),
)
p.drawText(
bbox.x(),
bbox.y(),
description,
)
pen = QtGui.QPen(QtGui.QColor("#000000"), 8, Qt.SolidLine)
p.setPen(pen)
for shape in self.shapes:
description = shape.description
if description:
bbox = shape.bounding_rect()
p.drawText(
bbox.x(),
bbox.y(),
description,
)
# if self.show_texts:
# p.setFont(
# QtGui.QFont(
# "Arial", int(max(6.0, int(round(8.0 / Shape.scale))))
# )
# )
# pen = QtGui.QPen(QtGui.QColor("#00FF00"), 8, Qt.SolidLine)
# p.setPen(pen)
# for shape in self.shapes:
# description = shape.description
# if description:
# #文字描述不能太长,否则显示会有问题
# if len(shape.description)>20:
# description=shape.description[0:20]
# bbox = shape.bounding_rect()
# fm = QtGui.QFontMetrics(p.font())
# rect = fm.boundingRect(description)
# print(bbox,rect)
# p.fillRect(
# rect.x() + bbox.x(),
# rect.y() + bbox.y(),
# rect.width(),
# rect.height(),
# QtGui.QColor("#00FF00"),
# )
# p.drawText(
# bbox.x(),
# bbox.y(),
# description,
# )
# pen = QtGui.QPen(QtGui.QColor("#000000"), 8, Qt.SolidLine)
# p.setPen(pen)
# for shape in self.shapes:
# description = shape.description
# if description:
# #文字描述不能太长,否则显示会有问题
# if len(shape.description)>20:
# description=shape.description[0:20]
# bbox = shape.bounding_rect()
# p.drawText(
# bbox.x(),
# bbox.y(),
# description,
# )

# Draw mouse coordinates
if self.show_cross_line:
Expand Down
3 changes: 2 additions & 1 deletion requirements-dev.txt
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
-r requirements.txt
build
twine
pyinstaller
pyinstaller
openai==0.27.2

0 comments on commit cf3ca4d

Please sign in to comment.